Implements exponential moving average shadowing for your model. Utilizes an inverse decay schedule to manage longer term training runs. By adjusting the power, you can control how fast EMA will ramp up to your specified beta. @crowsonkb's notes on EMA Warmup: If gamma=1 and p
| 30 | src.lerp_(tgt, weight) |
| 31 | |
| 32 | class EMA(Module): |
| 33 | """ |
| 34 | Implements exponential moving average shadowing for your model. |
| 35 | |
| 36 | Utilizes an inverse decay schedule to manage longer term training runs. |
| 37 | By adjusting the power, you can control how fast EMA will ramp up to your specified beta. |
| 38 | |
| 39 | @crowsonkb's notes on EMA Warmup: |
| 40 | |
| 41 | If gamma=1 and power=1, implements a simple average. gamma=1, power=2/3 are |
| 42 | good values for models you plan to train for a million or more steps (reaches decay |
| 43 | factor 0.999 at 31.6K steps, 0.9999 at 1M steps), gamma=1, power=3/4 for models |
| 44 | you plan to train for less (reaches decay factor 0.999 at 10K steps, 0.9999 at |
| 45 | 215.4k steps). |
| 46 | |
| 47 | Args: |
| 48 | inv_gamma (float): Inverse multiplicative factor of EMA warmup. Default: 1. |
| 49 | power (float): Exponential factor of EMA warmup. Default: 1. |
| 50 | min_value (float): The minimum EMA decay rate. Default: 0. |
| 51 | """ |
| 52 | |
| 53 | @beartype |
| 54 | def __init__( |
| 55 | self, |
| 56 | model: Module, |
| 57 | ema_model: Optional[Module] = None, # if your model has lazylinears or other types of non-deepcopyable modules, you can pass in your own ema model |
| 58 | beta = 0.9999, |
| 59 | update_after_step = 100, |
| 60 | update_every = 10, |
| 61 | inv_gamma = 1.0, |
| 62 | power = 2 / 3, |
| 63 | min_value = 0.0, |
| 64 | param_or_buffer_names_no_ema: Set[str] = set(), |
| 65 | ignore_names: Set[str] = set(), |
| 66 | ignore_startswith_names: Set[str] = set(), |
| 67 | include_online_model = True, # set this to False if you do not wish for the online model to be saved along with the ema model (managed externally) |
| 68 | allow_different_devices = False # if the EMA model is on a different device (say CPU), automatically move the tensor |
| 69 | ): |
| 70 | super().__init__() |
| 71 | self.beta = beta |
| 72 | |
| 73 | # whether to include the online model within the module tree, so that state_dict also saves it |
| 74 | |
| 75 | self.include_online_model = include_online_model |
| 76 | |
| 77 | if include_online_model: |
| 78 | self.online_model = model |
| 79 | else: |
| 80 | self.online_model = [model] # hack |
| 81 | |
| 82 | # ema model |
| 83 | |
| 84 | self.ema_model = ema_model |
| 85 | |
| 86 | if not exists(self.ema_model): |
| 87 | try: |
| 88 | self.ema_model = deepcopy(model) |
| 89 | except Exception as e: |