(
self,
model: Module,
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
beta = 0.9999,
update_after_step = 100,
update_every = 10,
inv_gamma = 1.0,
power = 2 / 3,
min_value = 0.0,
param_or_buffer_names_no_ema: Set[str] = set(),
ignore_names: Set[str] = set(),
ignore_startswith_names: Set[str] = set(),
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)
allow_different_devices = False # if the EMA model is on a different device (say CPU), automatically move the tensor
)
| 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: |
| 90 | print(f'Error: While trying to deepcopy model: {e}') |
| 91 | print('Your model was not copyable. Please make sure you are not using any LazyLinear') |
| 92 | exit() |
| 93 | |
| 94 | self.ema_model.requires_grad_(False) |
| 95 | |
| 96 | # parameter and buffer names |
| 97 | |
| 98 | self.parameter_names = {name for name, param in self.ema_model.named_parameters() if param.dtype in [torch.float, torch.float16]} |
| 99 | self.buffer_names = {name for name, buffer in self.ema_model.named_buffers() if buffer.dtype in [torch.float, torch.float16]} |
| 100 | |
| 101 | # tensor update functions |
| 102 | |
| 103 | self.inplace_copy = partial(inplace_copy, auto_move_device = allow_different_devices) |
| 104 | self.inplace_lerp = partial(inplace_lerp, auto_move_device = allow_different_devices) |
| 105 | |
| 106 | # updating hyperparameters |
| 107 | |
| 108 | self.update_every = update_every |
| 109 | self.update_after_step = update_after_step |
| 110 | |
| 111 | self.inv_gamma = inv_gamma |
nothing calls this directly
no test coverage detected