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 rest
(self, model)
| 75 | self.collected_params = [param.clone() for param in parameters] |
| 76 | |
| 77 | def restore(self, model): |
| 78 | """ |
| 79 | Restore the parameters stored with the `store` method. |
| 80 | Useful to validate the model with EMA parameters without affecting the |
| 81 | original optimization process. Store the parameters before the |
| 82 | `copy_to` method. After validation (or model saving), use this to |
| 83 | restore the former parameters. |
| 84 | Args: |
| 85 | model: A model that to restore its parameters. |
| 86 | """ |
| 87 | with GatheredParameters(model.parameters(), modifier_rank=0): |
| 88 | if deepspeed.comm.get_rank() == 0: |
| 89 | parameters = model.parameters() |
| 90 | for c_param, param in zip(self.collected_params, parameters): |
| 91 | param.data.copy_(c_param.data) |
| 92 | |
| 93 | @contextmanager |
| 94 | def activate(self, model): |