ema-pytorch

The EMAModuleWrapper is designed for complex, nested module trees (common in self-supervised learning). It automatically routes the outputs of target EMA submodules into the forward pass of specified online submodules.

How it works: When you call ema(x), the wrapper identifies the specified submodules and injects their corresponding EMA teacher outputs as keyword arguments (defaulting to ema_output) into the online module's forward method.

Configuration via ema_module_kwargs:

• Simple Mapping: {'online_path': 'target_path'}. The output of target_path is passed to online_path as ema_output.
• Advanced Mapping: {'online_path': {'ema_module_path': 'target_path', 'ema_kwarg': 'custom_name'}}. This allows you to specify a custom keyword argument name for the injected output.

Multi-view SSL: For scenarios where the student and teacher receive different augmented views, you can pass ema_args or ema_kwargs during the call: ema(student_input, ema_args = teacher_input).

from ema_pytorch import EMAModuleWrapper

# Example: Mapping branch_a.block1 to use branch_b.block1 as its teacher
ema = EMAModuleWrapper(
    model,
    beta = 0.99,
    ema_module_kwargs = {
        'branch_a.block1': 'branch_b.block1',
        'branch_a.block2': 'branch_b.block2'
    }
)

# Forwarding injects EMA outputs into the specified blocks
out, loss = ema(x)

# For multi-view SSL
out, loss = ema(student_input, ema_args = teacher_input)

The EMA class provides a simple way to maintain an Exponential Moving Average version of a PyTorch model.

Workflow:

1. Wrap your neural network with EMA(net, beta, ...).
2. Mutate your original network (net) using your optimizer (e.g., SGD).
3. Call ema.update() to update the moving average weights.
4. Use ema(data) to invoke the model with the EMA weights.

Key Parameters:

• net: Your PyTorch module.
• beta: The exponential moving average factor (decay).
• update_after_step: Number of .update() calls to wait before starting updates (warmup).
• update_every: Frequency of updates (e.g., 10 updates the EMA every 10th call to .update()).
• update_model_with_ema_every: (Optional) For 'Switch EMA' logic, specifies how often to sync the main model with the EMA model.

Saving and Accessing:

• It is recommended to save the entire EMA wrapper instance to preserve the step count and warmup logic.
• To access the EMA model directly, use ema.ema_model.

import torch
from ema_pytorch import EMA

net = torch.nn.Linear(512, 512)

ema = EMA(
    net,
    beta = 0.9999,
    update_after_step = 100,
    update_every = 10,
)

# Training loop simulation
for _ in range(10):
    # ... mutate net with optimizer ...
    ema.update()

# Inference with EMA weights
data = torch.randn(1, 512)
ema_output = ema(data)

The PostHocEMA class implements the method proposed by Karras et al. for synthesizing new EMA models from multiple existing ones.

Workflow:

1. Wrap your network with PostHocEMA.
2. Provide a tuple of sigma_rels (at least two values) to define the hyperparameters for multiple EMAs.
3. Specify a checkpoint_folder to store checkpoints for each sigma_rel across timesteps.
4. Call emas.update() during training.
5. Use emas.synthesize_ema_model(sigma_rel = value) to create a new synthesized model with a specific sigma_rel.

Key Parameters:

• sigma_rels: A tuple of hyperparameters for multiple EMAs.
• checkpoint_folder: Directory where checkpoints for each sigma_rel are saved.
• update_every: Frequency of updates.
• checkpoint_every_num_steps: Frequency of checkpointing.

import torch
from ema_pytorch import PostHocEMA

net = torch.nn.Linear(512, 512)

emas = PostHocEMA(
    net,
    sigma_rels = (0.05, 0.28),
    update_every = 10,
    checkpoint_every_num_steps = 10,
    checkpoint_folder = './post-hoc-ema-checkpoints'
)

# Training loop
for _ in range(1000):
    # ... mutate net ...
    emas.update()

# Synthesize a new model
synthesized_ema = emas.synthesize_ema_model(sigma_rel = 0.15)
output = synthesized_ema(torch.randn(1, 512))

The ema_module_kwargs parameter in EMAModuleWrapper defines how EMA outputs are harvested and injected. It supports several formats:

1. List/Set of paths: ['path.to.module']. The receiver and EMA path are identical, using the default_ema_kwarg (default: 'ema_output').
2. Dictionary mapping (String value): {'receiver_path': 'ema_path_or_kwarg'}.
   • If 'ema_path_or_kwarg' is a valid submodule path, it is treated as the ema_path and the default_ema_kwarg is used.
   • If it is NOT a valid path, it is treated as the kwarg_name to be used on the receiver_path (the EMA path defaults to the receiver path).
3. Dictionary mapping (Tuple/List value): {'receiver_path': (ema_path, kwarg_name, optional_transform)}.
4. Dictionary mapping (Dict value): {'receiver_path': {'ema_module_path': '...', 'ema_kwarg': '...', 'transform': func}}.

Note on transform: You can provide a callable to transform the EMA output before it is injected into the online module.

# 1. Simple list (receiver == ema_path)
ema_module_kwargs = ['blocks.0.attn']

# 2. Dict with string (mapping receiver to a specific kwarg name)
# Here 'blocks.0.attn' is treated as the kwarg name because it's not a valid submodule path
ema_module_kwargs = {'blocks.1.mlp': 'blocks.0.attn'}

# 3. Tuple specification
ema_module_kwargs = {'blocks.1.mlp': ('blocks.0.attn', 'ema_feat', my_transform_func)}

# 4. Full dict specification
ema_module_kwargs = {
    'blocks.1.mlp': {
        'ema_module_path': 'blocks.0.attn',
        'ema_kwarg': 'ema_feat',
        'transform': my_transform_func
    }
}