### Install ema-pytorch

Source: https://github.com/lucidrains/ema-pytorch/blob/main/README.md

Install the library using pip.

```bash
pip install ema-pytorch
```

--------------------------------

### Basic EMA Implementation

Source: https://context7.com/lucidrains/ema-pytorch/llms.txt

Wrap a PyTorch module with EMA to maintain averaged parameters with configurable decay and warmup.

```python
import torch
from ema_pytorch import EMA

# Create your neural network
net = torch.nn.Sequential(
    torch.nn.Linear(512, 256),
    torch.nn.ReLU(),
    torch.nn.Linear(256, 10)
)

# Wrap with EMA - specify decay (beta) and update parameters
ema = EMA(
    net,
    beta=0.9999,              # exponential moving average factor
    update_after_step=100,    # start updating after this many .update() calls
    update_every=10,          # how often to actually update (saves compute)
    inv_gamma=1.0,            # inverse multiplicative factor for EMA warmup
    power=2/3,                # exponential factor for EMA warmup
)

# Training loop
optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)

for step in range(1000):
    # Forward pass and optimization
    data = torch.randn(32, 512)
    target = torch.randint(0, 10, (32,))

    output = net(data)
    loss = torch.nn.functional.cross_entropy(output, target)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # Update EMA after each optimization step
    ema.update()

# Use EMA model for inference (same interface as original model)
data = torch.randn(1, 512)
ema_output = ema(data)  # calls ema.ema_model(data)

# Access the EMA model directly
print(ema.ema_model)

# Save the entire EMA wrapper (includes step count for warmup)
torch.save(ema.state_dict(), 'ema_checkpoint.pt')
```

--------------------------------

### Train with EMA on GPU

Source: https://context7.com/lucidrains/ema-pytorch/llms.txt

Demonstrates basic EMA usage during a training loop where the EMA model remains on the CPU while the online model is on the GPU.

```python
for step in range(1000):
    data = torch.randn(32, 512, device=device)
    output = net(data)

    with torch.no_grad():
        net.weight.add_(torch.randn_like(net.weight) * 0.01)

    ema.update()  # EMA stays on CPU

# For inference, EMA model will be moved to GPU if move_ema_to_online_device=True
ema_output = ema(torch.randn(1, 512, device=device))
```

--------------------------------

### Manual Checkpoint Control

Source: https://context7.com/lucidrains/ema-pytorch/llms.txt

Demonstrates how to manually trigger checkpoints for an EMA model after training epochs.

```python
for epoch in range(10):
    for step in range(100):
        with torch.no_grad():
            net.weight.add_(torch.randn_like(net.weight) * 0.01)
        emas.update()

    # Checkpoint at end of each epoch
    emas.checkpoint()
    print(f"Epoch {epoch}: Saved checkpoint at step {emas.step.item()}")

# Synthesize after training
best_ema = emas.synthesize_ema_model(sigma_rel=0.10)
```

--------------------------------

### Copying Parameters Between Models

Source: https://context7.com/lucidrains/ema-pytorch/llms.txt

Manually synchronizes weights between the online model and the EMA model.

```python
import torch
from ema_pytorch import EMA

net = torch.nn.Linear(512, 512)
ema = EMA(net, beta=0.9999)

# Initialize EMA with current model weights
ema.copy_params_from_model_to_ema()

# Train for a while...
for _ in range(1000):
    with torch.no_grad():
        net.weight.add_(torch.randn_like(net.weight) * 0.01)
    ema.update()

# Replace online model with EMA weights (useful for fine-tuning from EMA)
ema.copy_params_from_ema_to_model()

# Now net has the EMA weights
print(torch.allclose(net.weight, ema.ema_model.weight))  # True
```

--------------------------------

### Configure Manual Checkpointing for PostHocEMA

Source: https://context7.com/lucidrains/ema-pytorch/llms.txt

Disables automatic checkpointing in PostHocEMA to allow for manual control over when checkpoints are saved.

```python
import torch
from ema_pytorch import PostHocEMA

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

emas = PostHocEMA(
    net,
    sigma_rels=(0.05, 0.15, 0.28),
    update_every=10,
    checkpoint_every_num_steps='manual',  # disable automatic checkpointing
    checkpoint_folder='./manual-checkpoints',
)
```

--------------------------------

### Synthesize Models with PostHocEMA

Source: https://context7.com/lucidrains/ema-pytorch/llms.txt

Trains multiple EMA models and synthesizes a new one post-training using specific gamma or step values.

```python
import torch
from ema_pytorch import PostHocEMA

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

# Create PostHocEMA with multiple sigma_rels
emas = PostHocEMA(
    net,
    sigma_rels=(0.05, 0.28),  # need at least 2 different values
    update_every=10,
    checkpoint_every_num_steps=100,  # save checkpoints for synthesis
    checkpoint_folder='./post-hoc-ema-checkpoints',
    checkpoint_dtype=torch.float16,  # save checkpoints in fp16 to save space
)

# Training loop
net.train()
for step in range(1000):
    with torch.no_grad():
        net.weight.add_(torch.randn_like(net.weight) * 0.01)
        net.bias.add_(torch.randn_like(net.bias) * 0.01)

    emas.update()  # updates all EMA models and saves checkpoints

# After training: synthesize a new EMA with different sigma_rel
synthesized_ema = emas.synthesize_ema_model(sigma_rel=0.15)

# Use synthesized model for inference
data = torch.randn(1, 512)
output = synthesized_ema(data)

# Or synthesize with a specific gamma and step
synthesized_ema_v2 = emas.synthesize_ema_model(gamma=5.0, step=500)

# Get outputs from all EMA models at once
all_outputs = emas(data)  # returns tuple of outputs from each EMA model
```

--------------------------------

### EMA with Different Device Placement

Source: https://context7.com/lucidrains/ema-pytorch/llms.txt

Configure EMA to reside on a different device, such as CPU, while the model trains on GPU.

```python
import torch
from ema_pytorch import EMA

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Create model on GPU
net = torch.nn.Linear(512, 512).to(device)

# Create EMA model on CPU
ema_model = torch.nn.Linear(512, 512)  # stays on CPU

ema = EMA(
    net,
    ema_model=ema_model,
    beta=0.9999,
    allow_different_devices=True,  # allow EMA on different device
    move_ema_to_online_device=True,  # move EMA to GPU when needed for inference
)
```

--------------------------------

### Implement Karras EMA

Source: https://context7.com/lucidrains/ema-pytorch/llms.txt

Applies EMA using hyperparameters derived from Karras et al. diffusion model research.

```python
import torch
from ema_pytorch.post_hoc_ema import KarrasEMA

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

# Use sigma_rel (relative standard deviation) from the paper
karras_ema = KarrasEMA(
    net,
    sigma_rel=0.05,      # or use gamma directly: gamma=6.94
    update_every=100,    # update every 100 steps
    frozen=False,        # set True to freeze EMA weights
)

# Training loop
for step in range(10000):
    with torch.no_grad():
        net.weight.add_(torch.randn_like(net.weight) * 0.01)

    karras_ema.update()

    # The beta (decay) is computed dynamically based on step and gamma
    if step % 1000 == 0:
        print(f"Step {step}, beta: {karras_ema.beta:.6f}")

# Use for inference
output = karras_ema(torch.randn(1, 512))
```

--------------------------------

### Switch EMA for Continual Learning

Source: https://context7.com/lucidrains/ema-pytorch/llms.txt

Enable Switch EMA to periodically update the online model with EMA weights.

```python
import torch
from ema_pytorch import EMA

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

# Enable automatic model update from EMA every N steps
ema = EMA(
    net,
    beta=0.9999,
    update_after_step=100,
    update_every=10,
    update_model_with_ema_every=10000,  # update model weights from EMA every 10k steps
    update_model_with_ema_beta=0.0,     # amount of original model weight to keep (0 = full replacement)
)

# Training loop
for step in range(50000):
    with torch.no_grad():
        net.weight.copy_(torch.randn_like(net.weight))

    ema.update()  # Will automatically update model from EMA at step 10000, 20000, etc.

# Or manually trigger the update
ema.update_model_with_ema(decay=0.1)  # keep 10% of original model weights
```

--------------------------------

### Basic EMA Model Wrapping and Usage

Source: https://github.com/lucidrains/ema-pytorch/blob/main/README.md

Wrap your PyTorch neural network with EMA, specifying decay factor and update frequency. The EMA model can then be used similarly to the original network.

```python
import torch
from ema_pytorch import EMA

# your neural network as a pytorch module

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

# wrap your neural network, specify the decay (beta)

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 (updates every 10th .update() call)
)

# mutate your network, with SGD or otherwise

with torch.no_grad():
    net.weight.copy_(torch.randn_like(net.weight))
    net.bias.copy_(torch.randn_like(net.bias))

# you will call the update function on your moving average wrapper

ema.update()

# then, later on, you can invoke the EMA model the same way as your network

data = torch.randn(1, 512)

output     = net(data)	ema_output = ema(data)

# if you want to save your ema model, it is recommended you save the entire wrapper
# as it contains the number of steps taken (there is a warmup logic in there, recommended by @crowsonkb, validated for a number of projects now)
# however, if you wish to access the copy of your model with EMA, then it will live at ema.ema_model
```

--------------------------------

### Post-Hoc EMA Model Synthesis

Source: https://github.com/lucidrains/ema-pytorch/blob/main/README.md

Implement post-hoc synthesized EMA using Karras et al.'s method. Wrap your network and specify sigma_rels or gammas. After training, synthesize an EMA model with a desired sigma_rel.

```python
import torch
from ema_pytorch import PostHocEMA

# your neural network as a pytorch module

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

# wrap your neural network, specify the sigma_rels or gammas

emas = PostHocEMA(
    net,
    sigma_rels = (0.05, 0.28),           # a tuple with the hyperparameter for the multiple EMAs. you need at least 2 here to synthesize a new one
    update_every = 10,                  # how often to actually update, to save on compute (updates every 10th .update() call)
    checkpoint_every_num_steps = 10,
    checkpoint_folder = './post-hoc-ema-checkpoints'  # the folder of saved checkpoints for each sigma_rel (gamma) across timesteps with the hparam above, used to synthesizing a new EMA model after training
)

net.train()

for _ in range(1000):
    # mutate your network, with SGD or otherwise

    with torch.no_grad():
        net.weight.copy_(torch.randn_like(net.weight))
        net.bias.copy_(torch.randn_like(net.bias))

    # you will call the update function on your moving average wrapper

    emas.update()

# now that you have a few checkpoints
# you can synthesize an EMA model with a different sigma_rel (say 0.15)

synthesized_ema = emas.synthesize_ema_model(sigma_rel = 0.15)

# output with synthesized EMA

data = torch.randn(1, 512)

synthesized_ema_output = synthesized_ema(data)

```

--------------------------------

### EMA with Optimizer Hook

Source: https://context7.com/lucidrains/ema-pytorch/llms.txt

Register the EMA update as a post-step hook on the optimizer to automate updates.

```python
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)

optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)

# Register hook - update() is called automatically after optimizer.step()
hook_handle = ema.add_to_optimizer_post_step_hook(optimizer)

# Training loop without manual ema.update() calls
for step in range(1000):
    data = torch.randn(32, 512)
    target = torch.randn(32, 512)

    output = net(data)
    loss = torch.nn.functional.mse_loss(output, target)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()  # EMA automatically updated via hook

# Use EMA model for inference
ema_output = ema(torch.randn(1, 512))
```

--------------------------------

### Switch EMA Update Frequency

Source: https://github.com/lucidrains/ema-pytorch/blob/main/README.md

Configure the EMA wrapper to update the model with EMA weights at a specified frequency, or manually update at the end of an epoch.

```python
ema = EMA(
    net,
    ...,
    update_model_with_ema_every = 10000 # say 10k steps is 1 epoch
)

# or you can do it manually at the end of each epoch

ema.update_model_with_ema()
```

--------------------------------

### Apply EMA to Arbitrary Pytrees

Source: https://context7.com/lucidrains/ema-pytorch/llms.txt

Uses EMAPytree to track EMA for nested data structures like dictionaries and lists containing tensors.

```python
import torch
from ema_pytorch import EMAPytree

# Create an arbitrary pytree (dict with tensors)
online_tree = {
    'weights': torch.randn(512, 512),
    'bias': torch.randn(512),
    'nested': {
        'scale': torch.randn(1),
        'offset': torch.randn(1)
    }
}

# Wrap with EMAPytree
ema_tree = EMAPytree(
    online_tree,
    beta=0.9999,
    update_after_step=10,
    update_every=1,
)

# Simulate training updates
for step in range(100):
    with torch.no_grad():
        online_tree['weights'].add_(torch.randn_like(online_tree['weights']) * 0.01)
        online_tree['bias'].add_(torch.randn_like(online_tree['bias']) * 0.01)

    ema_tree.update()

# Access EMA pytree
print(ema_tree.ema_pytree['weights'])
print(ema_tree.ema_pytree['nested']['scale'])

# Get current decay value
print(f"Current decay: {ema_tree.get_current_decay()}")
```

--------------------------------

### EMA forward_eval Inference

Source: https://context7.com/lucidrains/ema-pytorch/llms.txt

Uses forward_eval to automatically handle eval mode and no_grad context during inference.

```python
import torch
from ema_pytorch import EMA

net = torch.nn.Sequential(
    torch.nn.Linear(512, 256),
    torch.nn.BatchNorm1d(256),  # behaves differently in train vs eval mode
    torch.nn.ReLU(),
    torch.nn.Linear(256, 10)
)

ema = EMA(net, beta=0.9999)

# Train for a bit
for _ in range(100):
    ema.update()

# forward_eval ensures no_grad and temporarily sets eval mode
data = torch.randn(32, 512)
output = ema.forward_eval(data)  # BatchNorm uses running stats, no gradients computed

# The EMA model's training state is preserved after forward_eval
print(f"EMA model training mode: {ema.ema_model.training}")
```

--------------------------------

### Exclude Parameters from EMA

Source: https://context7.com/lucidrains/ema-pytorch/llms.txt

Configures EMA to ignore or directly copy specific parameters or buffers using name-based filtering.

```python
import torch
from ema_pytorch import EMA

net = torch.nn.Sequential(
    torch.nn.Linear(512, 256),
    torch.nn.BatchNorm1d(256),
    torch.nn.Linear(256, 10)
)

ema = EMA(
    net,
    beta=0.9999,
    param_or_buffer_names_no_ema={'1.running_mean', '1.running_var'},  # copy directly, no EMA
    ignore_names={'1.num_batches_tracked'},  # completely ignore these
    ignore_startswith_names={'0.'},  # ignore all params starting with '0.'
)

for step in range(100):
    with torch.no_grad():
        for param in net.parameters():
            param.add_(torch.randn_like(param) * 0.01)
    ema.update()
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.