tfm.core.actions.EMACheckpointing

Eval action to save checkpoint with average weights when EMA is used.

tfm.core.actions.EMACheckpointing(
    export_dir: str,
    optimizer: tf.keras.optimizers.Optimizer,
    checkpoint: tf.train.Checkpoint,
    max_to_keep: int = 1
)

This action swaps the weights of the model with the average weights, then it saves the checkpoint under export_dir/ema_checkpoints. Checkpointing is expensive for large models, so doing this action in eval is more efficient than training.

export_dir str for the export directory of the EMA average weights.
optimizer tf.keras.optimizers.Optimizer optimizer instance used for training. This will be used to swap the model weights with the average weigths.
checkpoint tf.train.Checkpoint instance.
max_to_keep int for max checkpoints to keep in ema_checkpoints subdir.

Methods

__call__

View source

__call__(
    output: orbit.runner.Output
)

Swaps model weights, and saves the checkpoint.

Args
output The train or eval output.