aiaccel.torch.lightning.callbacks.NoBufferWeightAveraging

class aiaccel.torch.lightning.callbacks.NoBufferWeightAveraging(device: device | str | int | None = None, **kwargs: Any)[source]

Weight averaging callback that ignores buffers during averaging and swapping.

Checkpoints still store the current model buffers together with the averaged parameters so loading preserves non-averaged buffer state.

Example

>>> from lightning.pytorch import Trainer
>>> from torch.optim.swa_utils import get_ema_avg_fn
>>> ema = NoBufferWeightAveraging(avg_fn=get_ema_avg_fn(0.999))
>>> trainer = Trainer(callbacks=[ema])
__init__(device: device | str | int | None = None, **kwargs: Any)

Initialize the callback.

Parameters:

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

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

Methods

__init__([device])

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])

Called after every optimizer step and after every training epoch to check whether the average model should be updated.

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.