Switch Optimizer#

Module for the SwitchOptimizer callback.

class SwitchOptimizer(new_optimizers, epoch_switch)[source]#

Bases: Callback

Lightning callback for dynamically replacing optimizers during training.

This callback enables switching to one or more new optimizers at a specified epoch without restarting the training loop. It is particularly useful for staged optimization strategies (e.g., coarse-to-fine training or optimizer warm-up phases), where different optimizers are applied sequentially.

At the target epoch, the provided optimizers are hooked to the model parameters and replace the current optimizers in both the PINA solver and the Lightning trainer strategy.

Initialization of the SwitchOptimizer class.

Parameters:

  • new_optimizers (pina.optim.OptimizerInterface | list) – The model optimizers to switch to. Can be a single torch.optim.Optimizer instance or a list of them for multiple model solver.

  • epoch_switch (int) – The epoch at which the optimizer switch occurs.

Raises:

  • AssertionError – If epoch_switch is not a positive integer.

  • ValueError – If any of the provided optimizers are not instances of pina.optim.OptimizerInterface.

Example

>>> optimizer = TorchOptimizer(torch.optim.Adam, lr=0.01)
>>> switch_callback = SwitchOptimizer(
>>>     new_optimizers=optimizer, epoch_switch=10
>>> )
on_train_epoch_start(trainer, __)[source]#

Switch the optimizer at the start of the specified training epoch.

Parameters:

  • trainer (Trainer) – The trainer object managing the training process.

  • __ – Placeholder argument, not used.