kosmos.ml.config.factories.optimizer

Module Attributes

type ParamsT = Iterable[torch.Tensor] | Iterable[dict[str, Any]] | Iterable[tuple[str, torch.Tensor]]

Classes

class OptimizerConfig

Bases: abc.ABC

Optimizer configuration.

Methods

get_instance(params: ParamsT) → torch.optim.Optimizer

Get the optimizer instance.

Parameters:

params (ParamsT) – Parameters to optimize.

Returns:

Optimizer instance.

Return type:

Optimizer

---

class SGDOptimizerConfig(lr: float = 0.001, momentum: float = 0.0, weight_decay: float = 0.0, *, nesterov: bool = False)

Bases: OptimizerConfig

Stochastic gradient descent (SGD) optimizer configuration.

Initialize the SGD optimizer configuration.

Parameters:

• lr (float) – Learning rate. Defaults to 1e-3.

• momentum (float) – Momentum factor. Defaults to 0.0.

• weight_decay (float) – Weight decay. Defaults to 0.0.

• nesterov (bool) – Whether to use Nesterov momentum. Only applicable when momentum is non-zero. Defaults to False.

Methods

get_instance(params: ParamsT) → torch.optim.SGD

Get the SGD optimizer instance.

Parameters:

params (ParamsT) – Parameters to optimize.

Returns:

SGD optimizer instance.

Return type:

SGD

---

class AdamOptimizerConfig(lr: float = 0.001, weight_decay: float = 0.0)

Bases: OptimizerConfig

Adam optimizer configuration.

Initialize the Adam optimizer configuration.

Parameters:

• lr (float) – Learning rate. Defaults to 1e-3.

• weight_decay (float) – Weight decay. Defaults to 0.0.

Methods

get_instance(params: ParamsT) → torch.optim.Adam

Get the Adam optimizer instance.

Parameters:

params (ParamsT) – Parameters to optimize.

Returns:

Adam optimizer instance.

Return type:

Adam