Exponential moving averages of model parameters

environment.yml

@@ -17,3 +17,5 @@ dependencies:
   - pytest
   - psutil
   - gxx<12
+  - pip:
+    - torch-ema

torchmdnet/module.py

@@ -14,6 +14,7 @@
 from torchmdnet.models.model import create_model, load_model
 from torchmdnet.models.utils import dtype_mapping
 import torch_geometric.transforms as T
+from torch_ema import ExponentialMovingAverage
 
 
 class FloatCastDatasetWrapper(T.BaseTransform):
@@ -73,6 +74,11 @@ def __init__(self, hparams, prior_model=None, mean=None, std=None):
         else:
             self.model = create_model(self.hparams, prior_model, mean, std)
 
+        self.ema_prmtrs = None
+        if self.hparams.ema_prmtrs_decay is not None:
+            # initialize EMA for the model paremeters
+            self.ema_prmtrs = ExponentialMovingAverage(self.model.parameters(), decay=self.hparams.ema_prmtrs_decay)
+
         # initialize exponential smoothing
         self.ema = None
         self._reset_ema_dict()
@@ -251,6 +257,11 @@ def optimizer_step(self, *args, **kwargs):
                 pg["lr"] = lr_scale * self.hparams.lr
         super().optimizer_step(*args, **kwargs)
         optimizer.zero_grad()
+
+    def on_before_zero_grad(self, *args, **kwargs):
+        if self.ema_prmtrs is not None:
+            self.ema_prmtrs.to(self.device)
+            self.ema_prmtrs.update(self.model.parameters())
 
     def _get_mean_loss_dict_for_type(self, type):
         # Returns a list with the mean loss for each loss_fn for each stage (train, val, test)

torchmdnet/scripts/train.py

@@ -59,6 +59,7 @@ def get_argparse():
     parser.add_argument('--redirect', type=bool, default=False, help='Redirect stdout and stderr to log_dir/log')
     parser.add_argument('--gradient-clipping', type=float, default=0.0, help='Gradient clipping norm')
     parser.add_argument('--remove-ref-energy', action='store_true', help='If true, remove the reference energy from the dataset for delta-learning. Total energy can still be predicted by the model during inference by turning this flag off when loading.  The dataset must be compatible with Atomref for this to be used.')
+    parser.add_argument('--ema-prmtrs-decay', type=float, default=None, help='Exponential moving average decay for model parameters (None to disable)')
     # dataset specific
     parser.add_argument('--dataset', default=None, type=str, choices=datasets.__all__, help='Name of the torch_geometric dataset')
     parser.add_argument('--dataset-root', default='~/data', type=str, help='Data storage directory (not used if dataset is "CG")')