declare-lab/tango2

1import copy2import os3import random4from typing import Any, Dict, Iterable, Optional, Union5 6import numpy as np7import torch8 9from .utils import deprecate10 11 12def enable_full_determinism(seed: int):13    """14    Helper function for reproducible behavior during distributed training. See15    - https://pytorch.org/docs/stable/notes/randomness.html for pytorch16    """17    # set seed first18    set_seed(seed)19 20    #  Enable PyTorch deterministic mode. This potentially requires either the environment21    #  variable 'CUDA_LAUNCH_BLOCKING' or 'CUBLAS_WORKSPACE_CONFIG' to be set,22    # depending on the CUDA version, so we set them both here23    os.environ["CUDA_LAUNCH_BLOCKING"] = "1"24    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8"25    torch.use_deterministic_algorithms(True)26 27    # Enable CUDNN deterministic mode28    torch.backends.cudnn.deterministic = True29    torch.backends.cudnn.benchmark = False30 31 32def set_seed(seed: int):33    """34    Args:35    Helper function for reproducible behavior to set the seed in `random`, `numpy`, `torch`.36        seed (`int`): The seed to set.37    """38    random.seed(seed)39    np.random.seed(seed)40    torch.manual_seed(seed)41    torch.cuda.manual_seed_all(seed)42    # ^^ safe to call this function even if cuda is not available43 44 45# Adapted from torch-ema https://github.com/fadel/pytorch_ema/blob/master/torch_ema/ema.py#L1446class EMAModel:47    """48    Exponential Moving Average of models weights49    """50 51    def __init__(52        self,53        parameters: Iterable[torch.nn.Parameter],54        decay: float = 0.9999,55        min_decay: float = 0.0,56        update_after_step: int = 0,57        use_ema_warmup: bool = False,58        inv_gamma: Union[float, int] = 1.0,59        power: Union[float, int] = 2 / 3,60        model_cls: Optional[Any] = None,61        model_config: Dict[str, Any] = None,62        **kwargs,63    ):64        """65        Args:66            parameters (Iterable[torch.nn.Parameter]): The parameters to track.67            decay (float): The decay factor for the exponential moving average.68            min_decay (float): The minimum decay factor for the exponential moving average.69            update_after_step (int): The number of steps to wait before starting to update the EMA weights.70            use_ema_warmup (bool): Whether to use EMA warmup.71            inv_gamma (float):72                Inverse multiplicative factor of EMA warmup. Default: 1. Only used if `use_ema_warmup` is True.73            power (float): Exponential factor of EMA warmup. Default: 2/3. Only used if `use_ema_warmup` is True.74            device (Optional[Union[str, torch.device]]): The device to store the EMA weights on. If None, the EMA75                        weights will be stored on CPU.76 77        @crowsonkb's notes on EMA Warmup:78            If gamma=1 and power=1, implements a simple average. gamma=1, power=2/3 are good values for models you plan79            to train for a million or more steps (reaches decay factor 0.999 at 31.6K steps, 0.9999 at 1M steps),80            gamma=1, power=3/4 for models you plan to train for less (reaches decay factor 0.999 at 10K steps, 0.999981            at 215.4k steps).82        """83 84        if isinstance(parameters, torch.nn.Module):85            deprecation_message = (86                "Passing a `torch.nn.Module` to `ExponentialMovingAverage` is deprecated. "87                "Please pass the parameters of the module instead."88            )89            deprecate(90                "passing a `torch.nn.Module` to `ExponentialMovingAverage`",91                "1.0.0",92                deprecation_message,93                standard_warn=False,94            )95            parameters = parameters.parameters()96 97            # set use_ema_warmup to True if a torch.nn.Module is passed for backwards compatibility98            use_ema_warmup = True99 100        if kwargs.get("max_value", None) is not None:101            deprecation_message = "The `max_value` argument is deprecated. Please use `decay` instead."102            deprecate("max_value", "1.0.0", deprecation_message, standard_warn=False)103            decay = kwargs["max_value"]104 105        if kwargs.get("min_value", None) is not None:106            deprecation_message = "The `min_value` argument is deprecated. Please use `min_decay` instead."107            deprecate("min_value", "1.0.0", deprecation_message, standard_warn=False)108            min_decay = kwargs["min_value"]109 110        parameters = list(parameters)111        self.shadow_params = [p.clone().detach() for p in parameters]112 113        if kwargs.get("device", None) is not None:114            deprecation_message = "The `device` argument is deprecated. Please use `to` instead."115            deprecate("device", "1.0.0", deprecation_message, standard_warn=False)116            self.to(device=kwargs["device"])117 118        self.temp_stored_params = None119 120        self.decay = decay121        self.min_decay = min_decay122        self.update_after_step = update_after_step123        self.use_ema_warmup = use_ema_warmup124        self.inv_gamma = inv_gamma125        self.power = power126        self.optimization_step = 0127        self.cur_decay_value = None  # set in `step()`128 129        self.model_cls = model_cls130        self.model_config = model_config131 132    @classmethod133    def from_pretrained(cls, path, model_cls) -> "EMAModel":134        _, ema_kwargs = model_cls.load_config(path, return_unused_kwargs=True)135        model = model_cls.from_pretrained(path)136 137        ema_model = cls(model.parameters(), model_cls=model_cls, model_config=model.config)138 139        ema_model.load_state_dict(ema_kwargs)140        return ema_model141 142    def save_pretrained(self, path):143        if self.model_cls is None:144            raise ValueError("`save_pretrained` can only be used if `model_cls` was defined at __init__.")145 146        if self.model_config is None:147            raise ValueError("`save_pretrained` can only be used if `model_config` was defined at __init__.")148 149        model = self.model_cls.from_config(self.model_config)150        state_dict = self.state_dict()151        state_dict.pop("shadow_params", None)152 153        model.register_to_config(**state_dict)154        self.copy_to(model.parameters())155        model.save_pretrained(path)156 157    def get_decay(self, optimization_step: int) -> float:158        """159        Compute the decay factor for the exponential moving average.160        """161        step = max(0, optimization_step - self.update_after_step - 1)162 163        if step <= 0:164            return 0.0165 166        if self.use_ema_warmup:167            cur_decay_value = 1 - (1 + step / self.inv_gamma) ** -self.power168        else:169            cur_decay_value = (1 + step) / (10 + step)170 171        cur_decay_value = min(cur_decay_value, self.decay)172        # make sure decay is not smaller than min_decay173        cur_decay_value = max(cur_decay_value, self.min_decay)174        return cur_decay_value175 176    @torch.no_grad()177    def step(self, parameters: Iterable[torch.nn.Parameter]):178        if isinstance(parameters, torch.nn.Module):179            deprecation_message = (180                "Passing a `torch.nn.Module` to `ExponentialMovingAverage.step` is deprecated. "181                "Please pass the parameters of the module instead."182            )183            deprecate(184                "passing a `torch.nn.Module` to `ExponentialMovingAverage.step`",185                "1.0.0",186                deprecation_message,187                standard_warn=False,188            )189            parameters = parameters.parameters()190 191        parameters = list(parameters)192 193        self.optimization_step += 1194 195        # Compute the decay factor for the exponential moving average.196        decay = self.get_decay(self.optimization_step)197        self.cur_decay_value = decay198        one_minus_decay = 1 - decay199 200        for s_param, param in zip(self.shadow_params, parameters):201            if param.requires_grad:202                s_param.sub_(one_minus_decay * (s_param - param))203            else:204                s_param.copy_(param)205 206    def copy_to(self, parameters: Iterable[torch.nn.Parameter]) -> None:207        """208        Copy current averaged parameters into given collection of parameters.209 210        Args:211            parameters: Iterable of `torch.nn.Parameter`; the parameters to be212                updated with the stored moving averages. If `None`, the parameters with which this213                `ExponentialMovingAverage` was initialized will be used.214        """215        parameters = list(parameters)216        for s_param, param in zip(self.shadow_params, parameters):217            param.data.copy_(s_param.to(param.device).data)218 219    def to(self, device=None, dtype=None) -> None:220        r"""Move internal buffers of the ExponentialMovingAverage to `device`.221 222        Args:223            device: like `device` argument to `torch.Tensor.to`224        """225        # .to() on the tensors handles None correctly226        self.shadow_params = [227            p.to(device=device, dtype=dtype) if p.is_floating_point() else p.to(device=device)228            for p in self.shadow_params229        ]230 231    def state_dict(self) -> dict:232        r"""233        Returns the state of the ExponentialMovingAverage as a dict. This method is used by accelerate during234        checkpointing to save the ema state dict.235        """236        # Following PyTorch conventions, references to tensors are returned:237        # "returns a reference to the state and not its copy!" -238        # https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict239        return {240            "decay": self.decay,241            "min_decay": self.min_decay,242            "optimization_step": self.optimization_step,243            "update_after_step": self.update_after_step,244            "use_ema_warmup": self.use_ema_warmup,245            "inv_gamma": self.inv_gamma,246            "power": self.power,247            "shadow_params": self.shadow_params,248        }249 250    def store(self, parameters: Iterable[torch.nn.Parameter]) -> None:251        r"""252        Args:253        Save the current parameters for restoring later.254            parameters: Iterable of `torch.nn.Parameter`; the parameters to be255                temporarily stored.256        """257        self.temp_stored_params = [param.detach().cpu().clone() for param in parameters]258 259    def restore(self, parameters: Iterable[torch.nn.Parameter]) -> None:260        r"""261        Args:262        Restore the parameters stored with the `store` method. Useful to validate the model with EMA parameters without:263        affecting the original optimization process. Store the parameters before the `copy_to()` method. After264        validation (or model saving), use this to restore the former parameters.265            parameters: Iterable of `torch.nn.Parameter`; the parameters to be266                updated with the stored parameters. If `None`, the parameters with which this267                `ExponentialMovingAverage` was initialized will be used.268        """269        if self.temp_stored_params is None:270            raise RuntimeError("This ExponentialMovingAverage has no `store()`ed weights " "to `restore()`")271        for c_param, param in zip(self.temp_stored_params, parameters):272            param.data.copy_(c_param.data)273 274        # Better memory-wise.275        self.temp_stored_params = None276 277    def load_state_dict(self, state_dict: dict) -> None:278        r"""279        Args:280        Loads the ExponentialMovingAverage state. This method is used by accelerate during checkpointing to save the281        ema state dict.282            state_dict (dict): EMA state. Should be an object returned283                from a call to :meth:`state_dict`.284        """285        # deepcopy, to be consistent with module API286        state_dict = copy.deepcopy(state_dict)287 288        self.decay = state_dict.get("decay", self.decay)289        if self.decay < 0.0 or self.decay > 1.0:290            raise ValueError("Decay must be between 0 and 1")291 292        self.min_decay = state_dict.get("min_decay", self.min_decay)293        if not isinstance(self.min_decay, float):294            raise ValueError("Invalid min_decay")295 296        self.optimization_step = state_dict.get("optimization_step", self.optimization_step)297        if not isinstance(self.optimization_step, int):298            raise ValueError("Invalid optimization_step")299 300        self.update_after_step = state_dict.get("update_after_step", self.update_after_step)301        if not isinstance(self.update_after_step, int):302            raise ValueError("Invalid update_after_step")303 304        self.use_ema_warmup = state_dict.get("use_ema_warmup", self.use_ema_warmup)305        if not isinstance(self.use_ema_warmup, bool):306            raise ValueError("Invalid use_ema_warmup")307 308        self.inv_gamma = state_dict.get("inv_gamma", self.inv_gamma)309        if not isinstance(self.inv_gamma, (float, int)):310            raise ValueError("Invalid inv_gamma")311 312        self.power = state_dict.get("power", self.power)313        if not isinstance(self.power, (float, int)):314            raise ValueError("Invalid power")315 316        shadow_params = state_dict.get("shadow_params", None)317        if shadow_params is not None:318            self.shadow_params = shadow_params319            if not isinstance(self.shadow_params, list):320                raise ValueError("shadow_params must be a list")321            if not all(isinstance(p, torch.Tensor) for p in self.shadow_params):322                raise ValueError("shadow_params must all be Tensors")323