Optimizers

Any torch optimizer or custom optimizer can be used for training.

Tip

To create custom optimizers, refer to the custom optimizers tutorial.

Torch Optimizers

Torch optimizers (torch.optim) can be specified as relative python import paths for the _target_ attribute in the configuration file. Any additional attributes (except id) are passed as keyword arguments to the optimizer constructor.

For example, torch.optim.Adam can be used as follows:

conf/optimizer/Adam.yaml

id: Adam
_target_: torch.optim.Adam

autrainer provides a number of default configurations for torch optimizers:

Default Configurations

conf/optimizer/SGD.yaml

id: SGD
_target_: torch.optim.SGD

conf/optimizer/Adam.yaml

id: Adam
_target_: torch.optim.Adam

conf/optimizer/AdamW.yaml

id: AdamW
_target_: torch.optim.AdamW

conf/optimizer/Adadelta.yaml

id: Adadelta
_target_: torch.optim.Adadelta

conf/optimizer/Adagrad.yaml

id: Adagrad
_target_: torch.optim.Adagrad

conf/optimizer/Adamax.yaml

id: Adamax
_target_: torch.optim.Adamax

conf/optimizer/NAdam.yaml

id: NAdam
_target_: torch.optim.NAdam

conf/optimizer/RAdam.yaml

id: RAdam
_target_: torch.optim.RAdam

conf/optimizer/SparseAdam.yaml

id: SparseAdam
_target_: torch.optim.SparseAdam

conf/optimizer/RMSprop.yaml

id: RMSprop
_target_: torch.optim.RMSprop

conf/optimizer/Rprop.yaml

id: Rprop
_target_: torch.optim.Rprop

conf/optimizer/ASGD.yaml

id: ASGD
_target_: torch.optim.ASGD

conf/optimizer/LBFGS.yaml

id: LBFGS
_target_: torch.optim.LBFGS

Custom Optimizers

class autrainer.optimizers.SAM(params, base_optimizer, rho=0.05, adaptive=False, **kwargs)

Sharpness Aware Minimization (SAM) optimizer.

This implementation is adapted from the following repository: davda54/sam

For more information, see: https://arxiv.org/abs/2010.01412

Parameters:

• params (Module) – Model parameters.

• base_optimizer (str) – Underlying optimizer performing the sharpness-aware update, specified as a relative import path.

• rho (float) – Size of the neighborhood for computing the max loss. Defaults to 0.05.

• adaptive (bool) – Whether to use an experimental implementation of element-wise Adaptive SAM. Defaults to False.

• **kwargs – Additional arguments passed to the underlying optimizer.

Default Configurations

conf/optimizer/SAM-SGD.yaml

id: SAM-SGD
_target_: autrainer.optimizers.SAM
base_optimizer: torch.optim.SGD

Custom Step Function

Custom optimizers can optionally provide a custom_step() function that is called instead of the standard training step and should be defined as follows:

custom_step function of an optimizer

class SomeOptimizer(torch.optim.Optimizer):
    def custom_step(
        self,
        model: AbstractModel,
        data: DataBatch,
        criterion: torch.nn.Module,
        probabilities_fn: Callable,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Custom step function for the optimizer.

        Args:
            model: The model to train.
            data: The data batch containing features, target, and potentially
                additional fields. The data batch is on the same
                device as the model. Additional fields are passed to the model
                as keyword arguments if they are present in the model's forward
                method.
            criterion: Loss function.
            probabilities_fn: Function to convert model outputs to
                probabilities.

        Returns:
            Tuple containing the non-reduced loss and model outputs.
        """

Note

The custom_step() function should perform both the forward and backward pass as well as update the model parameters accordingly.