ema-pytorch

library0.7.9pypypiunverified

ema-pytorch is a Python library that provides an easy way to integrate Exponential Moving Average (EMA) into PyTorch models. It helps stabilize training and improve generalization by maintaining a smoothed version of model parameters over time. The library is actively developed, with frequent updates, and is currently at version 0.7.9.

ai-ml

pip install ema-pytorch

github ↗package ↗

INSTALL

IMPORT

SIG · EMA-PYTORCH

ema-pytorch

ai-mlpythonv0.7.9

Install

66.6s avg

Import

—

Disk

4787MB

Pass rate

4/ 10

Env Coverage4 / 10

glibc

3.9–3.13

musl

3.9–3.13

Install & Compatibility

Where this runs

tested against v0.7.9 · pip install

no network on importno background threads

Install × environment matrix

Each cell = how many times install + import succeeded across repeated harness runs. Partial = flaky.
glibc = Debian/Ubuntu slim · musl = Alpine Linux

musl

glibc

py 3.10

✕ build_error

✓ 76.1s

py 3.11

✕ build_error

✓ 69.2s

py 3.12

✕ build_error

✓ 63.98s

py 3.13

✕ build_error

✓ 57.3s

py 3.9

✕ build_error

✕ timeout

4787MB installed

● package 4787MB

Code

Verified usage

Verified import paths — ran on the pinned version, not inferred.

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.

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}")

Debug

Known issues

gotchaWhen 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.

affects: All

gotchaFor 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.

affects: All

gotchaExtended 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).

affects: All

gotchaTo 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`).

affects: All

Upgrade

Version history

0.7.9latest on PyPI · released Dec 19, 2025

Audit

Dependencies

torchrequiredCore deep learning framework for model operations.

Agent activity

20 hits · last 30 days

node

OpenAI (training)

Resources

githubhttps://github.com/lucidrains/ema-pytorch ↗

packagehttps://pypi.org/project/ema-pytorch/ ↗