EMA PyTorch

Common errors

• Error: While trying to deepcopy model: {e} Your model was not copyable. Please make sure you are not using any LazyLinear

  cause The EMA wrapper attempts to deepcopy the model, and certain PyTorch modules like `torch.nn.LazyLinear` cannot be deepcopied until their input shape is determined.
  fix Ensure that your model's layers are fully initialized with concrete input/output shapes before passing it to `EMA`. Avoid `LazyLinear` if you intend to deepcopy, or ensure it's evaluated once before EMA initialization. The `PostHocEMA` source code specifically mentions `LazyLinear` as an issue.

• KeyError: 'ema_model' not found when loading checkpoint

  cause You might be trying to load only the `state_dict` of the EMA model (`ema.ema_model.state_dict()`) into a new `EMA` wrapper, or you're trying to load a general `state_dict` without the specific keys `ema-pytorch` expects.
  fix As recommended, save the entire `ema` object using `torch.save(ema, 'path/to/ema.pth')`. When loading, use `ema = torch.load('path/to/ema.pth')`. If you must save `state_dict`, ensure all relevant components are saved (model, ema_model, and the EMA wrapper's internal state) and loaded carefully.

• AttributeError: 'EMA' object has no attribute 'train' or 'eval'

  cause You are trying to call `.train()` or `.eval()` directly on the `ema` wrapper object, but these methods should be called on the *wrapped* model, or the `ema.ema_model` directly.
  fix To set the underlying EMA model to evaluation mode, use `ema.ema_model.eval()`. Similarly, for training mode (though `ema_model` is typically for inference), use `ema.ema_model.train()` or ensure the main `net` is in `train()` mode for updates. The `ema` wrapper itself handles its state.

Warnings

• gotcha When using EMA with optimizers that employ weight decay (e.g., AdamW), there can be interference. EMA tracks raw weights, not their decayed counterparts, which might lead to the EMA not fully accounting for the optimizer's weight decay impact.
  Fix: Consider using a lower weight decay rate when EMA is active, or apply EMA updates only to the optimizer's non-decayed weights through a custom schedule.
• gotcha For Mixed Precision Training (e.g., using `torch.cuda.amp.autocast`), EMA updates should occur *outside* the `autocast` context to prevent numerical instabilities and precision issues.
  Fix: Ensure `ema.update()` is called after `scaler.update()` and outside the `with autocast():` block.
• gotcha Extended training can sometimes destabilize EMA decay, especially with decay rates close to 1, potentially leading to 'oversmoothing'.
  Fix: Experiment with resetting EMA decay to a lower rate after certain epochs and gradually increasing it, or implementing a decay rate schedule where it increases as the model converges (e.g., from 0.99 to 0.999 over epochs).
• gotcha To correctly save and load the EMA state, it's recommended to save the entire EMA wrapper object, not just `ema.ema_model.state_dict()`, as the wrapper contains crucial state like the number of steps taken (for warmup logic).
  Fix: Use `torch.save(ema, 'ema_model.pth')` and `ema = torch.load('ema_model.pth')`. If saving only the state_dict, ensure you also save `ema.num_updates` and `ema.beta` (and potentially `update_after_step`).

Install

• pip install ema-pytorch Install stable version

Imports

• EMA

  from ema_pytorch import EMA

  Primary class for standard Exponential Moving Average.
• PostHocEMA

  from ema_pytorch import PostHocEMA

  Class for post-hoc EMA synthesis, proposed by Karras et al.

Quickstart

Initialize your PyTorch model, then wrap it with the `EMA` class, specifying the decay factor (`beta`). During your training loop, call `ema.update()` after `optimizer.step()` to update the EMA parameters. For inference or validation, you can directly call the `ema` object, which will use the averaged parameters.

import torch
from ema_pytorch import EMA

# Your neural network as a PyTorch module
net = torch.nn.Linear(512, 512)

# Wrap your neural network with EMA
ema = EMA(
    net,
    beta = 0.9999,  # exponential moving average factor
    update_after_step = 100, # only after this number of .update() calls will it start updating
    update_every = 10  # how often to actually update, to save on compute
)

# Simulate training steps
optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)
for step in range(1000):
    optimizer.zero_grad()
    data = torch.randn(1, 512)
    target = torch.randn(1, 512)
    output = net(data)
    loss = torch.nn.functional.mse_loss(output, target)
    loss.backward()
    optimizer.step()
    
    # Update the EMA model
    ema.update()

# Later, for inference, use the EMA model
with torch.no_grad():
    data_inference = torch.randn(1, 512)
    ema_output = ema(data_inference)
    print(f"EMA model output shape: {ema_output.shape}")