blob: 09f3e18295ab3bc72be13939743102dffbf1f855 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
|
# pytorch_ema
A very small library for computing exponential moving averages of model
parameters.
This library was written for personal use. Nevertheless, if you run into issues
or have suggestions for improvement, feel free to open either a new issue or
pull request.
## Installation
```
pip install -U git+https://github.com/fadel/pytorch_ema
```
## Example
```python
import torch
import torch.nn.functional as F
from torch_ema import ExponentialMovingAverage
x_train = torch.rand((100, 10))
y_train = torch.rand(100).round().long()
model = torch.nn.Linear(10, 2)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
ema = ExponentialMovingAverage(model.parameters(), decay=0.995)
# Train for a few epochs
model.train()
for _ in range(10):
logits = model(x_train)
loss = F.cross_entropy(logits, y_train)
optimizer.zero_grad()
loss.backward()
optimizer.step()
ema.update(model.parameters())
# Compare losses:
# Original
model.eval()
logits = model(x_train)
loss = F.cross_entropy(logits, y_train)
print(loss.item())
# With EMA
ema.copy_to(model.parameters())
logits = model(x_train)
loss = F.cross_entropy(logits, y_train)
print(loss.item())
```
|
