| 89 | |
| 90 | |
| 91 | class EMA: |
| 92 | # Found this useful (thanks alexis-jacq): |
| 93 | # https://discuss.pytorch.org/t/how-to-apply-exponential-moving-average-decay-for-variables/10856/3 |
| 94 | def __init__(self, gamma=0.99, save=True, save_frequency=100, save_filename="ema_weights.pth"): |
| 95 | """ |
| 96 | Initialize the weight to which we will do the |
| 97 | exponential moving average and the dictionary |
| 98 | where we store the model parameters |
| 99 | """ |
| 100 | self.gamma = gamma |
| 101 | self.registered = {} |
| 102 | self.save_filename = save_filename |
| 103 | self.save_frequency = save_frequency |
| 104 | self.count = 0 |
| 105 | |
| 106 | if save_filename in os.listdir("."): |
| 107 | self.registered = torch.load(self.save_filename) |
| 108 | |
| 109 | if not save: |
| 110 | warnings.warn("Note that the exponential moving average weights will not be saved to a .pth file!") |
| 111 | |
| 112 | def register_weights(self, model): |
| 113 | """ |
| 114 | Registers the weights of the model which will |
| 115 | later be used when we take the moving average |
| 116 | """ |
| 117 | for name, param in model.named_parameters(): |
| 118 | if param.requires_grad: |
| 119 | self.registered[name] = param.clone().detach() |
| 120 | |
| 121 | def __call__(self, model): |
| 122 | self.count += 1 |
| 123 | for name, param in model.named_parameters(): |
| 124 | if param.requires_grad: |
| 125 | new_weight = param.clone().detach() if name not in self.registered else self.gamma * param + (1 - self.gamma) * self.registered[name] |
| 126 | self.registered[name] = new_weight |
| 127 | |
| 128 | if self.count % self.save_frequency == 0: |
| 129 | self.save_ema_weights() |
| 130 | |
| 131 | def copy_weights_to(self, model): |
| 132 | for name, param in model.named_parameters(): |
| 133 | if param.requires_grad: |
| 134 | param.data = self.registered[name] |
| 135 | |
| 136 | def save_ema_weights(self): |
| 137 | torch.save(self.registered, self.save_filename) |