r""" Args: Restore the parameters stored with the `store` method. Useful to validate the model with EMA parameters without: affecting the original optimization process. Store the parameters before the `copy_to()` method. After validation (or model saving), use this to
(self, parameters: Iterable[torch.nn.Parameter])
| 536 | self.temp_stored_params = [param.detach().cpu().clone() for param in parameters] |
| 537 | |
| 538 | def restore(self, parameters: Iterable[torch.nn.Parameter]) -> None: |
| 539 | r""" |
| 540 | Args: |
| 541 | Restore the parameters stored with the `store` method. Useful to validate the model with EMA parameters without: |
| 542 | affecting the original optimization process. Store the parameters before the `copy_to()` method. After |
| 543 | validation (or model saving), use this to restore the former parameters. |
| 544 | parameters: Iterable of `torch.nn.Parameter`; the parameters to be |
| 545 | updated with the stored parameters. If `None`, the parameters with which this |
| 546 | `ExponentialMovingAverage` was initialized will be used. |
| 547 | """ |
| 548 | if self.temp_stored_params is None: |
| 549 | raise RuntimeError("This ExponentialMovingAverage has no `store()`ed weights " "to `restore()`") |
| 550 | if self.foreach: |
| 551 | torch._foreach_copy_( |
| 552 | [param.data for param in parameters], [c_param.data for c_param in self.temp_stored_params] |
| 553 | ) |
| 554 | else: |
| 555 | for c_param, param in zip(self.temp_stored_params, parameters): |
| 556 | param.data.copy_(c_param.data) |
| 557 | |
| 558 | # Better memory-wise. |
| 559 | self.temp_stored_params = None |
| 560 | |
| 561 | def load_state_dict(self, state_dict: dict) -> None: |
| 562 | r""" |