kd.optim.UseEmaParams#

class kauldron.optim.UseEmaParams( *, ema_params_transform: str | None = None, partial_ok: bool = False, )[source]

Bases: kauldron.checkpoints.partial_loader.AbstractPartialLoader

Use the EMA parameters stored by the ema_params transform.

ema_params_transform

The path to the ema_params transform in the optax.chain or optax.named_chain. If not set, the state is searched for an EmaParamsState (which has to be unique).

Type:: str | None

partial_ok

If True, missing EMA params are ignored and the original params are kept. This is useful in combination with frozen parameters, e.g. when using with kd.optim.partial_updates.

Type:: bool

ema_params_transform: str | None = None

partial_ok: bool = False

transform(state)[source]: Replace the parameters with the params from the EMA state.