Add version hyperparameter to the checkpoints.

torchmdnet/__init__.py

@@ -0,0 +1,15 @@
+import importlib.metadata
+import subprocess
+
+try:
+    __version__ = importlib.metadata.version("torchmd-net")
+except importlib.metadata.PackageNotFoundError:
+    try:
+        __version__ = (
+            subprocess.check_output(["git", "describe", "--abbrev=0", "--tags"])
+            .strip()
+            .decode("utf-8")
+        )
+    except:
+        print("Failed to retrieve the current version, defaulting to 0")
+        __version__ = "0"

torchmdnet/models/model.py

@@ -7,13 +7,14 @@
 import torch
 from torch.autograd import grad
 from torch import nn, Tensor
+import torchmdnet
 from torchmdnet.models import output_modules
 from torchmdnet.models.wrappers import AtomFilter
 from torchmdnet.models.utils import dtype_mapping
 from torchmdnet import priors
 from lightning_utilities.core.rank_zero import rank_zero_warn
 import warnings
-
+from packaging import version
 
 def create_model(args, prior_model=None, mean=None, std=None):
     """Create a model from the given arguments.
@@ -156,6 +157,13 @@ def load_model(filepath, args=None, device="cpu", **kwargs):
     if args is None:
         args = ckpt["hyper_parameters"]
 
+    # Check the version that the checkpoint was created with
+    ckpt_version = ckpt.get("version", "0.15.2") # Default to the first version that introduced the version key
+    current_version = torchmdnet.__version__
+    if version.parse(ckpt_version) != version.parse(current_version):
+        warnings.warn(
+            f"Checkpoint was created with version {version}, current version is {current_version}."
+        )
     delta_learning = args["remove_ref_energy"] if "remove_ref_energy" in args else False
 
     for key, value in kwargs.items():

torchmdnet/module.py

@@ -9,8 +9,9 @@
 from torch.nn.functional import local_response_norm, mse_loss, l1_loss
 from torch import Tensor
 from typing import Optional, Dict, Tuple
-
+import warnings
 from lightning import LightningModule
+import torchmdnet
 from torchmdnet.models.model import create_model, load_model
 from torchmdnet.models.utils import dtype_mapping
 import torch_geometric.transforms as T
@@ -65,7 +66,8 @@ def __init__(self, hparams, prior_model=None, mean=None, std=None):
             hparams["charge"] = False
         if "spin" not in hparams:
             hparams["spin"] = False
-
+        if "version" not in hparams:
+            hparams["version"] = torchmdnet.__version__
         self.save_hyperparameters(hparams)
 
         if self.hparams.load_model: