aiaccel.torch.lightning.callbacks.NoBufferEMAWeightAveraging

class aiaccel.torch.lightning.callbacks.NoBufferEMAWeightAveraging(device: device | str | int | None = None, decay: float = 0.999, update_every_n_steps: int = 1, update_starting_at_step: int | None = None, update_starting_at_epoch: int | None = None, **kwargs: Any)[source]

Exponential moving average (EMA) callback that ignores buffers.

__init__(device: device | str | int | None = None, decay: float = 0.999, update_every_n_steps: int = 1, update_starting_at_step: int | None = None, update_starting_at_epoch: int | None = None, **kwargs: Any)

Initialize the callback.

Parameters:

  • device – Device that stores the averaged model. If None, the current model device is used.

  • decay – Decay factor for the exponential moving average. Should be between 0 and 1. Default is 0.999.

  • update_every_n_steps – Update EMA every N optimizer steps.

  • update_starting_at_step – Start EMA updates at or after this optimizer step.

  • update_starting_at_epoch – Start EMA updates at or after this epoch.

  • **kwargs – Additional arguments forwarded to lightning.pytorch.callbacks.EMAWeightAveraging. use_buffers is always fixed to False in this subclass.

Methods

__init__([device, decay, ...])

Initialize the callback.

load_state_dict(state_dict)

Called when loading a checkpoint.

on_after_backward(trainer, pl_module)

Called after loss.backward() and before optimizers are stepped.

on_before_backward(trainer, pl_module, loss)

Called before loss.backward().

on_before_optimizer_step(trainer, pl_module, ...)

Called before optimizer.step().

on_before_zero_grad(trainer, pl_module, ...)

Called before optimizer.zero_grad().

on_exception(trainer, pl_module, exception)

Called when any trainer execution is interrupted by an exception.

on_fit_end(trainer, pl_module)

Called when fit ends.

on_fit_start(trainer, pl_module)

Called when fit begins.

on_load_checkpoint(trainer, pl_module, ...)

Called when loading a model checkpoint.

on_predict_batch_end(trainer, pl_module, ...)

Called when the predict batch ends.

on_predict_batch_start(trainer, pl_module, ...)

Called when the predict batch begins.

on_predict_end(trainer, pl_module)

Called when predict ends.

on_predict_epoch_end(trainer, pl_module)

Called when the predict epoch ends.

on_predict_epoch_start(trainer, pl_module)

Called when the predict epoch begins.

on_predict_start(trainer, pl_module)

Called when the predict begins.

on_sanity_check_end(trainer, pl_module)

Called when the validation sanity check ends.

on_sanity_check_start(trainer, pl_module)

Called when the validation sanity check starts.

on_save_checkpoint(trainer, pl_module, ...)

Called when saving a checkpoint.

on_test_batch_end(trainer, pl_module, ...[, ...])

Called when the test batch ends.

on_test_batch_start(trainer, pl_module, ...)

Called when the test batch begins.

on_test_end(trainer, pl_module)

Called when the test ends.

on_test_epoch_end(trainer, pl_module)

Called when the test epoch ends.

on_test_epoch_start(trainer, pl_module)

Called when the test epoch begins.

on_test_start(trainer, pl_module)

Called when the test begins.

on_train_batch_end(trainer, pl_module, ...)

Called when a training batch ends.

on_train_batch_start(trainer, pl_module, ...)

Called when the train batch begins.

on_train_end(trainer, pl_module)

Called when training ends.

on_train_epoch_end(trainer, pl_module)

Called when a training epoch ends.

on_train_epoch_start(trainer, pl_module)

Called when the train epoch begins.

on_train_start(trainer, pl_module)

Called when the train begins.

on_validation_batch_end(trainer, pl_module, ...)

Called when the validation batch ends.

on_validation_batch_start(trainer, ...[, ...])

Called when the validation batch begins.

on_validation_end(trainer, pl_module)

Called when the validation loop ends.

on_validation_epoch_end(trainer, pl_module)

Called when a validation epoch ends.

on_validation_epoch_start(trainer, pl_module)

Called when a validation epoch begins.

on_validation_start(trainer, pl_module)

Called when the validation loop begins.

setup(trainer, pl_module, stage)

Called when fit, validate, test, predict, or tune begins.

should_update([step_idx, epoch_idx])

Decide when to update the model weights.

state_dict()

Called when saving a checkpoint.

teardown(trainer, pl_module, stage)

Called when fit, validate, test, predict, or tune ends.

Attributes

state_key

Identifier for the state of the callback.