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

Adam

conf/optimizer/Adam.yaml

id: Adam
_target_: torch.optim.Adam

conf/optimizer/Adamax.yaml

id: Adamax
_target_: torch.optim.Adamax

conf/optimizer/AdamW.yaml

id: AdamW
_target_: torch.optim.AdamW

SGD

conf/optimizer/SGD.yaml

id: SGD
_target_: torch.optim.SGD

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: torch.nn.Module,  # model
        data: torch.Tensor,  # batched input data
        target: torch.Tensor,  # batched target data
        criterion: torch.nn.Module,  # loss function
    ) -> Tuple[float, torch.Tensor]:
        loss = ...  # reduced loss over the batch
        outputs = ...  # detached model outputs
        return loss, outputs

Note

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