EMA¶

def __call__(self, model: BaseModel | nn.parallel.DistributedDataParallel):
    """
    Update the moving average of the parameters.

    Does not mutate the input, only updates the moving average.

    Args:
        model: The model whose parameters should be updated. Should be a model
            specified identically to the one passed when this object was
            instantiated.
    """
    decay = self.decay

    self.num_updates += 1
    if self._faster_decay_at_start:
        decay = torch.min(
            self.decay, (1 + self.num_updates) / (10 + self.num_updates)
        )
    self.cur_decay = decay
    with torch.no_grad():
        module_parameters = dict(model.named_parameters())

        for key in module_parameters:
            ema_name = self._get_ema_name(key)
            if module_parameters[key].requires_grad:
                self._ema_params[ema_name] = self._ema_params[ema_name].type_as(
                    module_parameters[key]
                )
                # EMA_params = decay * EMA_params + (1 - decay) * model_params
                self._ema_params[ema_name].sub_(
                    (1.0 - decay)
                    * (self._ema_params[ema_name] - module_parameters[key])
                )
            elif ema_name in self._ema_params:
                raise ValueError(
                    f"Expected model parameter {key} to require gradient, "
                    "but it does not"
                )

def copy_to(self, model: BaseModel | nn.parallel.DistributedDataParallel):
    """
    Copy the averaged parameters to the model, overwriting its values.
    """
    m_param = dict(model.named_parameters())
    for key in m_param:
        if m_param[key].requires_grad:
            m_param[key].data.copy_(self._ema_params[self._get_ema_name(key)].data)
        else:
            assert self._get_ema_name(key) not in self._ema_params

def store(self, parameters: Iterable[nn.Parameter]):
    """
    Save the current parameters for restoring later.

    Args:
        parameters: The parameters to be stored for later restoration by `restore`
    """
    self._stored_params = [param.clone() for param in parameters]

def restore(self, parameters: Iterable[nn.Parameter]):
    """
    Restore the parameters stored with the `store` method.

    Useful to validate the model with EMA parameters without affecting the
    original optimization process. Store the parameters before the
    `copy_to` method. After validation (or model saving), use this to
    restore the former parameters.

    Args:
        parameters: The parameters to be updated with the values stored by `store`
    """
    for c_param, param in zip(self._stored_params, parameters, strict=True):
        param.data.copy_(c_param.data)

    self._stored_params = []

def get_state(self, include_ema_params: bool):
    """
    Get the state of the EMA tracker.

    Args:
        include_ema_params: Whether to include the EMA parameters in the state.
        You probably want this to be True when saving a checkpoint meant to resume
        training (i.e. not an EMA checkpoint) or False when saving a checkpoint for
        evaluation (i.e. a checkpoint where the model weights are the EMA weights).

    Returns:
        The state of the EMA tracker.
    """
    state = {
        "decay": self.decay,
        "num_updates": self.num_updates,
        "faster_decay_at_start": self._faster_decay_at_start,
    }
    if include_ema_params:
        state["ema_params"] = self._ema_params
    return state

@classmethod
def from_state(cls, state, model: torch.nn.Module) -> "EMATracker":
    """
    Create an EMA tracker from a state.

    Args:
        state: The state of the EMA tracker.
        model: The model whose parameters should be tracked, used to
            initialize the EMA weights if they are not provided in the state.

    Returns:
        The EMA tracker.
    """
    ema = cls(model, float(state["decay"]), state["faster_decay_at_start"])
    ema.num_updates = state["num_updates"]
    if ema_params := state.get("ema_params"):
        unexpected_keys = ema_params.keys() - ema._ema_params.keys()
        missing_keys = ema._ema_params.keys() - ema_params.keys()
        assert not unexpected_keys and not missing_keys, (
            f"EMA parameters keys do not match. "
            f"This is likely due to a mismatch between the model and checkpoint. "
            f"Unexpected keys: {unexpected_keys}"
        )
        ema._ema_params = ema_params
    return ema