Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models Keep a moving average of everything in the model state_dict (parameters and buffers). This is intended to allow functionality like https://www.tensorflow.org/api_docs/python/tf/train/ExponentialM
| 28 | |
| 29 | |
| 30 | class ModelEMA: |
| 31 | """ |
| 32 | Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models |
| 33 | Keep a moving average of everything in the model state_dict (parameters and buffers). |
| 34 | This is intended to allow functionality like |
| 35 | https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage |
| 36 | A smoothed version of the weights is necessary for some training schemes to perform well. |
| 37 | This class is sensitive where it is initialized in the sequence of model init, |
| 38 | GPU assignment and distributed training wrappers. |
| 39 | """ |
| 40 | |
| 41 | def __init__(self, model, decay=0.9999, updates=0): |
| 42 | """ |
| 43 | Args: |
| 44 | model (nn.Module): model to apply EMA. |
| 45 | decay (float): ema decay reate. |
| 46 | updates (int): counter of EMA updates. |
| 47 | """ |
| 48 | # Create EMA(FP32) |
| 49 | self.ema = deepcopy(model.module if is_parallel(model) else model).eval() |
| 50 | self.updates = updates |
| 51 | # decay exponential ramp (to help early epochs) |
| 52 | self.decay = lambda x: decay * (1 - math.exp(-x / 2000)) |
| 53 | for p in self.ema.parameters(): |
| 54 | p.requires_grad_(False) |
| 55 | |
| 56 | def update(self, model): |
| 57 | # Update EMA parameters |
| 58 | with torch.no_grad(): |
| 59 | self.updates += 1 |
| 60 | d = self.decay(self.updates) |
| 61 | |
| 62 | msd = ( |
| 63 | model.module.state_dict() if is_parallel(model) else model.state_dict() |
| 64 | ) # model state_dict |
| 65 | for k, v in self.ema.state_dict().items(): |
| 66 | if v.dtype.is_floating_point: |
| 67 | v *= d |
| 68 | v += (1.0 - d) * msd[k].detach() |
| 69 | |
| 70 | def update_attr(self, model, include=(), exclude=("process_group", "reducer")): |
| 71 | # Update EMA attributes |
| 72 | copy_attr(self.ema, model, include, exclude) |