espnet3.components.callbacks.default_callbacks.AverageCheckpointsCallback

source

class espnet3.components.callbacks.default_callbacks.AverageCheckpointsCallback(output_dir, best_ckpt_callbacks)

Bases: Callback

A custom callback for weight averaging over the top-K checkpoints.

This can be useful to smooth out fluctuations in weights across the best-performing models and can lead to improved generalization performance at inference time.

Behavior. : - Loads the state_dict from each of the top-K checkpoints saved by given ModelCheckpoint callbacks.

• Averages every parameter in the checkpoint’s model state_dict. If the keys match the current model’s own state_dict() keys, they are averaged as-is; otherwise a model. prefix is stripped from each key before checking again.
• Ignores or simply accumulates integer-type parameters (e.g., BatchNorm’s num_batches_tracked).
• Saves the averaged model as a .pth file in output_dir.

• Parameters:
  • output_dir (str or Path) – The directory where the averaged model will be saved.
  • best_ckpt_callbacks (List *[*ModelCheckpoint ]) – A list of ModelCheckpoint callbacks whose top-K checkpoints will be used for averaging. Each callback must have best_k_models populated.

Notes

• The final filename will be: : {monitor_name}.ave_{K}best.pth
• This callback only runs on the global rank 0 process : (for distributed training).

• Raises:KeyError – If checkpoint keys are inconsistent across the top-K checkpoints, or if the averaged keys (with or without a model. prefix) do not match the current model’s state_dict() keys.

Example

>>> avg_ckpt_cb = AverageCheckpointsCallback(
...     output_dir="checkpoints/",
...     best_ckpt_callbacks=[val_loss_ckpt_cb, acc_ckpt_cb]
... )
>>> trainer = Trainer(callbacks=[avg_ckpt_cb])

Initialize AverageCheckpointsCallback object.

on_validation_end(trainer, pl_module)

At the end of validation, average the top-K checkpoints and save.