| from collections import OrderedDict |
|
|
| import torch |
|
|
|
|
| @torch.no_grad() |
| def update_ema( |
| ema_model: torch.nn.Module, model: torch.nn.Module, optimizer=None, decay: float = 0.9999, sharded: bool = True |
| ) -> None: |
| """ |
| Step the EMA model towards the current model. |
| """ |
| ema_params = OrderedDict(ema_model.named_parameters()) |
| model_params = OrderedDict(model.named_parameters()) |
|
|
| for name, param in model_params.items(): |
| if name == "pos_embed": |
| continue |
| if param.requires_grad == False: |
| continue |
| if not sharded: |
| param_data = param.data |
| ema_params[name].mul_(decay).add_(param_data, alpha=1 - decay) |
| else: |
| if param.data.dtype != torch.float32: |
| param_id = id(param) |
| master_param = optimizer._param_store.working_to_master_param[param_id] |
| param_data = master_param.data |
| else: |
| param_data = param.data |
| ema_params[name].mul_(decay).add_(param_data, alpha=1 - decay) |
|
|
