(self, iteration)
| 42 | |
| 43 | @torch.no_grad() |
| 44 | def update(self, iteration): |
| 45 | self.iteration = iteration |
| 46 | if iteration > self.start_iter: |
| 47 | for p, p_ema in zip( |
| 48 | self.model.model.parameters(), |
| 49 | self.ema_model.model.parameters(), |
| 50 | ): |
| 51 | p_ema.data = p_ema.data * self.decay + \ |
| 52 | p.data * (1 - self.decay) |
| 53 | else: |
| 54 | for p, p_ema in zip( |
| 55 | self.model.model.parameters(), |
| 56 | self.ema_model.model.parameters(), |
| 57 | ): |
| 58 | p_ema.data = p.data.clone().detach() |
| 59 | |
| 60 | def load_ema_params(self): |
| 61 | # load ema params to model |