Initialize the weight to which we will do the exponential moving average and the dictionary where we store the model parameters
(self, gamma=0.99, save=True, save_frequency=100, save_filename="ema_weights.pth")
| 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 | """ |
nothing calls this directly
no outgoing calls
no test coverage detected