From 448588ac2401e9b68d5fb3f660d38d00cf633aa1 Mon Sep 17 00:00:00 2001 From: zehui-lin Date: Fri, 15 Jan 2021 11:29:57 +0800 Subject: Add Feature: restore --- torch_ema/ema.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/torch_ema/ema.py b/torch_ema/ema.py index 4e9dd99..20c62c4 100644 --- a/torch_ema/ema.py +++ b/torch_ema/ema.py @@ -54,6 +54,22 @@ class ExponentialMovingAverage: parameters: Iterable of `torch.nn.Parameter`; the parameters to be updated with the stored moving averages. """ + self.collected_parameters = [] for s_param, param in zip(self.shadow_params, parameters): + self.collected_parameters.append(param.clone()) if param.requires_grad: param.data.copy_(s_param.data) + + def restore(self, parameters): + """ + Restore the parameters given to the copy_to function. + Usually used in validation. Want to validate the model with EMA parameters without affecting the original optimization process. + + Args: + parameters: Iterable of `torch.nn.Parameter`; the parameters to be + updated with the stored original parameters. + """ + for s_param, param in zip(self.collected_parameters, parameters): + if param.requires_grad: + param.data.copy_(s_param.data) + -- cgit v1.2.3