add a way to disable the EMA updating

so can basically disable EMA with a flag.
This is another difference between u-net trained on score_sde side deterministically
and the separate deterministic training approach.""

In theory decay rate of 1 should allow this but it's complicated by a num_updates params too

src/ml_downscaling_emulator/score_sde_pytorch/configs/deterministic/default_configs.py

@@ -71,6 +71,7 @@ def get_default_configs():
   model.loc_spec_channels = 0
   model.num_scales = 1
   model.ema_rate = 0.9999
+  model.ema_disabled = False
   model.dropout = 0.1
   model.embedding_type = 'fourier'
 

src/ml_downscaling_emulator/score_sde_pytorch/configs/deterministic/ukcp_local_pr_12em_plain_unet.py

@@ -43,6 +43,7 @@ def get_config():
   # model
   model = config.model
   model.name = 'det_cunet'
+  model.ema_disabled = True
 
   # optimizer
   optim = config.optim

src/ml_downscaling_emulator/score_sde_pytorch/models/ema.py

@@ -12,7 +12,7 @@ class ExponentialMovingAverage:
   Maintains (exponential) moving average of a set of parameters.
   """
 
-  def __init__(self, parameters, decay, use_num_updates=True):
+  def __init__(self, parameters, decay, use_num_updates=True, disable_update=False):
     """
     Args:
       parameters: Iterable of `torch.nn.Parameter`; usually the result of
@@ -28,6 +28,7 @@ def __init__(self, parameters, decay, use_num_updates=True):
     self.shadow_params = [p.clone().detach()
                           for p in parameters if p.requires_grad]
     self.collected_params = []
+    self.disable_update = disable_update
 
   def update(self, parameters):
     """
@@ -40,15 +41,19 @@ def update(self, parameters):
       parameters: Iterable of `torch.nn.Parameter`; usually the same set of
         parameters used to initialize this object.
     """
-    decay = self.decay
-    if self.num_updates is not None:
-      self.num_updates += 1
-      decay = min(decay, (1 + self.num_updates) / (10 + self.num_updates))
-    one_minus_decay = 1.0 - decay
-    with torch.no_grad():
-      parameters = [p for p in parameters if p.requires_grad]
-      for s_param, param in zip(self.shadow_params, parameters):
-        s_param.sub_(one_minus_decay * (s_param - param))
+    if not self.disable_update:
+      decay = self.decay
+      if self.num_updates is not None:
+        self.num_updates += 1
+        decay = min(decay, (1 + self.num_updates) / (10 + self.num_updates))
+      one_minus_decay = 1.0 - decay
+      with torch.no_grad():
+        parameters = [p for p in parameters if p.requires_grad]
+        for s_param, param in zip(self.shadow_params, parameters):
+          s_param.sub_(one_minus_decay * (s_param - param))
+    else:
+      # if disabled then just maintain a copy of the parameters
+      self.shadow_params = [p.clone().detach() for p in parameters if p.requires_grad]
 
   def copy_to(self, parameters):
     """

src/ml_downscaling_emulator/score_sde_pytorch/run_lib.py

@@ -124,7 +124,8 @@ def train(config, workdir):
     location_params = LocationParams(config.model.loc_spec_channels, config.data.image_size)
     location_params = location_params.to(config.device)
     location_params = torch.nn.DataParallel(location_params)
-    ema = ExponentialMovingAverage(itertools.chain(score_model.parameters(), location_params.parameters()), decay=config.model.ema_rate)
+    ema = ExponentialMovingAverage(itertools.chain(score_model.parameters(), location_params.parameters()), decay=config.model.ema_rate, disable_update=config.model.ema_disabled)
+
     optimizer = losses.get_optimizer(config, itertools.chain(score_model.parameters(), location_params.parameters()))
     state = dict(optimizer=optimizer, model=score_model, location_params=location_params, ema=ema, step=0, epoch=0)