---
id: ema-pytorch
version: "0.8.3"
license: MIT
license_treatment: permissive
maintenance: active
---

# ema-pytorch — Easy way to keep track of exponential moving average version of your pytorch module

License: permissive · Maintenance: active · Downloads: 253.1K/mo

## What it is and what it does

ema-pytorch wraps a PyTorch neural network to automatically maintain an exponential moving average (EMA) copy of its weights during training. Instead of using the latest model weights for inference or validation, you can use the EMA version, which often generalizes better by smoothing out noisy weight updates. The wrapper handles the decay schedule, update frequency, and optional warmup logic so you only need to call update() after each training step.

The package supports three main patterns: standard EMA with configurable decay and update frequency, post-hoc EMA synthesis (where multiple EMA versions are checkpointed and interpolated to create a new EMA at a different decay rate), and target representation routing for nested module architectures where EMA outputs are injected as teacher signals into student submodules. It is commonly used in diffusion models, self-supervised learning, and other training regimes where model averaging improves convergence or generalization.

Use it for:

- Improve validation accuracy in diffusion models by using EMA weights instead of the latest training weights.
- Implement teacher-student training in self-supervised learning where EMA outputs guide online module predictions.
- Synthesize a new EMA model at a different decay rate from saved checkpoints without retraining.
- Reduce overfitting in standard supervised learning by averaging weights across training steps.
- Experiment with Switch EMA to update the online model with EMA weights periodically for better loss landscape flatness.

## Worth the install?

AI-flagged interpretation of the facts above — verify before relying.

Maintains an exponential moving average (EMA) copy of a PyTorch model during training, with configurable decay, update frequency, and optional post-hoc synthesis across multiple EMA versions.

Yes. The package is lightweight, actively maintained, has no known vulnerabilities, and solves a well-defined problem (EMA tracking) that is standard in modern deep learning. Install it if you are training a PyTorch model and want to leverage weight averaging for better generalization or if you need post-hoc EMA synthesis for research.

## Install

pip install ema-pytorch
uv add ema-pytorch
poetry add ema-pytorch

## Installing ema-pytorch

Before you install: Low friction: pure Python wheel with only torch as a runtime dependency. Active maintenance with a recent release (14 days old) and steady commit history.

License in practice: MIT license permits commercial and private use with minimal restrictions; you must include a copy of the license and copyright notice.

Quickstart:

pip install ema-pytorch

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)

# After training updates to net:
ema.update()
output = ema(data)

Requires PyTorch to be installed; Python 3.8 or later.

Verify before relying:

- Whether post-hoc EMA synthesis (Karras et al.) offers measurable accuracy gains for your model architecture.
- Performance overhead of EMA updates relative to model size and batch frequency.
- Compatibility with distributed training frameworks (DDP, FSDP) not mentioned in the description.

## Package facts

- License: MIT (permissive)
- Python support: supports_current
- Install friction: low
- Maintenance: active
- Downloads: 253.1K/month (top 15,000 on PyPI)
- Known vulnerabilities: none known

## Tags

exponential moving average pytorch, EMA model wrapper, pytorch model averaging, teacher student EMA, model weight smoothing, training dynamics EMA, post-hoc EMA synthesis, model-training, weight-averaging, diffusion-models

[View on SkillFed](https://skillfed.io/packages/ema-pytorch) · [View on PyPI](https://pypi.org/project/ema-pytorch/)