class_docstr: Exponential Moving Average of Fairseq Models
class_docstr_tok: ["Exponential","Moving","Average","of","Fairseq","Models"]
docstr: One update of the EMA model based on new model weights
docstr_tok: ["One","update","of","the","EMA","model","based","on","new","model","weights"]
code: def _step_internal(self, new_model):
      """One update of the EMA model based on new model weights"""
      decay = self.decay
  
      ema_state_dict = {}
      ema_params = (
          self.fp32_params
          if self.config.ema_fp32
          else self.model.state_dict()
      )
  
      new_p = []
      ema_p = []
  
      for key, param in new_model.named_parameters():
          if isinstance(param, dict):
              continue
  
          if not self.add_missing_params and key not in ema_params:
              continue
  
          try:
              ema_param = ema_params[key]
          except KeyError:
              ema_param = (
                  param.float().clone()
                  if param.ndim == 1
                  else copy.deepcopy(param)
              )
              ema_params[key] = ema_param
  
          if param.shape != ema_param.shape:
              raise ValueError(
                  "incompatible tensor shapes between model param and ema param"
                  + "{} vs. {}".format(param.shape, ema_param.shape)
              )
  
          if "version" in key:
              # Do not decay a model.version pytorch param
              continue
  
          lr = 1 - decay
  
          if key in self.skip_keys or not param.requires_grad:
              ema_params[key].copy_(
                  param.to(dtype=ema_param.dtype).data
              )
              ema_param = ema_params[key]
          else:
              if self.log_norms:
                  new_p.append(param)
                  ema_p.append(ema_param)
  
              ema_param.mul_(1 - lr)
              ema_param.add_(
                  param.data.to(dtype=ema_param.dtype), alpha=lr
              )
  
          ema_state_dict[key] = ema_param
  
      for key, param in new_model.named_buffers():
          ema_state_dict[key] = param
  
      if self.log_norms:
          if "model_norm" in self.logs:
              self.prev_model_norm = self.logs["model_norm"]
  
          chunk_size = 2048 * 32
          has_inf = torch.zeros(
              (1, 1),
              dtype=torch.int,
              device=next(new_model.parameters()).device,
          )
  
          new_norm = multi_tensor_l2norm(
              chunk_size, has_inf, [new_p], False
          )
          old_norm = multi_tensor_l2norm(
              chunk_size, has_inf, [ema_p], False
          )
  
          self.logs["model_norm"] = new_norm[0]
          self.logs["ema_norm"] = old_norm[0]
  
      self.restore(ema_state_dict, build_fp32_params=False)
  
file_path: fairseq/fairseq/modules/ema_module.py
imports: import copy
  from dataclasses import dataclass, field
  import logging
  import torch
  from omegaconf import II
  from fairseq.dataclass import FairseqDataclass
label_desc: Use image node_id 6 for calling the EMAModule obj's underlying member method code with example usage: obj._step_internal(new_model) without return types