Step the EMA model towards the current model.
(ema_model, model, decay=0.9999)
| 102 | |
| 103 | @torch.no_grad() |
| 104 | def update_ema(ema_model, model, decay=0.9999): |
| 105 | """ |
| 106 | Step the EMA model towards the current model. |
| 107 | """ |
| 108 | ema_params = OrderedDict(ema_model.named_parameters()) |
| 109 | model_params = OrderedDict(model.named_parameters()) |
| 110 | |
| 111 | for name, param in model_params.items(): |
| 112 | # TODO: Consider applying only to params that require_grad to avoid small numerical changes of pos_embed |
| 113 | ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay) |
| 114 | |
| 115 | |
| 116 | def requires_grad(model, flag=True): |