Spaces:
Running on Zero
Running on Zero
File size: 3,361 Bytes
ff7b988 | 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 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | import torch
from torch._utils import _flatten_dense_tensors
import numpy as np
# EMA always in float, as accumulation needs lots of bits
class EMA:
def __init__(self, params, mu=0.999):
self.mu = mu
self.state = [(p, self.get_model_state(p)) for p in params if p.requires_grad]
def get_model_state(self, p):
return p.data.float().detach().clone()
def step(self):
for p, state in self.state:
state.mul_(self.mu).add_(1 - self.mu, p.data.float())
def swap(self):
# swap ema and model params
for p, state in self.state:
other_state = self.get_model_state(p)
p.data.copy_(state.type_as(p.data))
state.copy_(other_state)
class CPUEMA:
def __init__(self, params, mu=0.999, freq=1):
self.mu = mu**freq
self.state = [(p, self.get_model_state(p)) for p in params if p.requires_grad]
self.freq = freq
self.steps = 0
def get_model_state(self, p):
with torch.no_grad():
state = p.data.float().detach().cpu().numpy()
return state
def step(self):
with torch.no_grad():
self.steps += 1
if self.steps % self.freq == 0:
for i in range(len(self.state)):
p, state = self.state[i]
state = torch.from_numpy(state).cuda()
state.mul_(self.mu).add_(1 - self.mu, p.data.float())
self.state[i] = (p, state.cpu().numpy())
def swap(self):
with torch.no_grad():
# swap ema and model params
for p, state in self.state:
other_state = self.get_model_state(p)
p.data.copy_(torch.from_numpy(state).type_as(p.data))
np.copyto(state, other_state)
class FusedEMA:
def __init__(self, params, mu=0.999):
self.mu = mu
params = list(params)
self.params = {}
self.params['fp16'] = [p for p in params if p.requires_grad and p.data.dtype == torch.float16]
self.params['fp32'] = [p for p in params if p.requires_grad and p.data.dtype != torch.float16]
self.groups = [group for group in self.params.keys() if len(self.params[group]) > 0]
self.state = {}
for group in self.groups:
self.state[group] = self.get_model_state(group)
def get_model_state(self, group):
params = self.params[group]
return _flatten_dense_tensors([p.data.float() for p in params])
# if self.fp16:
# return _flatten_dense_tensors([p.data.half() for p in self.param_group if p.dtype])
# else:
# return _flatten_dense_tensors([p.data for p in self.param_group])
def step(self):
for group in self.groups:
self.state[group].mul_(self.mu).add_(1 - self.mu, self.get_model_state(group))
def swap(self):
# swap ema and model params
for group in self.groups:
other_state = self.get_model_state(group)
state = self.state[group]
params = self.params[group]
offset = 0
for p in params:
numel = p.data.numel()
p.data = state.narrow(0, offset, numel).view_as(p.data).type_as(p.data)
offset += numel
self.state[group] = other_state
|