Step the EMA model towards the current model.
(ema_model, model, decay=0.9999)
| 75 | |
| 76 | @torch.no_grad() |
| 77 | def update_ema(ema_model, model, decay=0.9999): |
| 78 | """ |
| 79 | Step the EMA model towards the current model. |
| 80 | """ |
| 81 | ema_params = OrderedDict(ema_model.named_parameters()) |
| 82 | model_params = OrderedDict(model.named_parameters()) |
| 83 | assert set(ema_params.keys()) == set(model_params.keys()) |
| 84 | |
| 85 | for name, param in model_params.items(): |
| 86 | # TODO: Consider applying only to params that require_grad to avoid small numerical changes of pos_embed |
| 87 | ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay) |
| 88 | |
| 89 | |
| 90 | def cleanup(): |