Repository: chemprop/chemprop
Pull request #1191: Make converting 2.0 to 2.1 models easier and more obvious
Author: login: JacksonBurns
Created: 2025-03-13T17:25:19Z | Merged: 2025-05-18T02:11:41Z

## Description
As seen in https://github.com/chemprop/chemprop/issues/1156 when a user tries to load a chemprop 2.0 model in 2.1 they will get an `AttributeError`.
Solution is not obvious.

## Example / Current workflow
User would have to read the release notes from 2.1.0 to find the announcement about the conversion script to fix this error.

## Bugfix / Desired workflow
Expose conversion script through the CLI, suggest using it when the corresponding error is raised on an attempted load.

## Questions
We could do this other ways - just first idea.

I've not tested this, as I don't have any checkpoints on hand.


Diff:
diff --git a/chemprop/cli/convert.py b/chemprop/cli/convert.py
index c8558c06c..a47796383 100644
--- a/chemprop/cli/convert.py
+++ b/chemprop/cli/convert.py
@@ -5,44 +5,60 @@
 
 from chemprop.cli.utils import Subcommand
 from chemprop.utils.v1_to_v2 import convert_model_file_v1_to_v2
+from chemprop.utils.v2_0_to_v2_1 import convert_model_file_v2_0_to_v2_1
 
 logger = logging.getLogger(__name__)
 
 
 class ConvertSubcommand(Subcommand):
     COMMAND = "convert"
-    HELP = "Convert a v1 model checkpoint (.pt) to a v2 model checkpoint (.pt)."
+    HELP = "Convert model checkpoint (.pt) to more recent version (.pt)."
 
     @classmethod
     def add_args(cls, parser: ArgumentParser) -> ArgumentParser:
+        parser.add_argument(
+            "-c",
+            "--conversion",
+            choices=["v1_to_v2", "v2_0_to_v2_1"],
+            help="Conversion to perform. Models converted from v1 to v2 must be run with the v1 featurizer via `--multi-hot-atom-featurizer-mode v1`.",
+            default="v1_to_v2",
+        )
         parser.add_argument(
             "-i",
             "--input-path",
             required=True,
             type=Path,
-            help="Path to a v1 model .pt checkpoint file",
+            help="Path to a model .pt checkpoint file",
         )
         parser.add_argument(
             "-o",
             "--output-path",
             type=Path,
-            help="Path to which the converted model will be saved (``CURRENT_DIRECTORY/STEM_OF_INPUT_v2.pt`` by default)",
+            help="Path to which the converted model will be saved (``CURRENT_DIRECTORY/STEM_OF_INPUT_newversion.pt`` by default)",
         )
         return parser
 
     @classmethod
     def func(cls, args: Namespace):
+        _suffix = None
+        if args.conversion == "v1_to_v2":
+            _suffix = "_v2.pt"
+        elif args.conversion == "v2_0_to_v2_1":
+            _suffix = "_v2_1.pt"
         if args.output_path is None:
-            args.output_path = Path(args.input_path.stem + "_v2.pt")
+            args.output_path = Path(args.input_path.stem + _suffix)
         if args.output_path.suffix != ".pt":
             raise ArgumentError(
                 argument=None, message=f"Output must be a `.pt` file. Got {args.output_path}"
             )
 
         logger.info(
-            f"Converting v1 model checkpoint '{args.input_path}' to v2 model checkpoint '{args.output_path}'..."
+            f"Converting model checkpoint '{args.input_path}' to model checkpoint '{args.output_path}'..."
         )
-        convert_model_file_v1_to_v2(args.input_path, args.output_path)
+        if args.conversion == "v1_to_v2":
+            convert_model_file_v1_to_v2(args.input_path, args.output_path)
+        elif args.conversion == "v2_0_to_v2_1":
+            convert_model_file_v2_0_to_v2_1(args.input_path, args.output_path)
 
 
 if __name__ == "__main__":
diff --git a/chemprop/models/model.py b/chemprop/models/model.py
index b527b4444..1cc059c17 100644
--- a/chemprop/models/model.py
+++ b/chemprop/models/model.py
@@ -2,6 +2,7 @@
 
 import io
 import logging
+import traceback
 from typing import Iterable, TypeAlias
 
 from lightning import pytorch as pl
@@ -253,7 +254,12 @@ def get_batch_size(self, batch: TrainingBatch) -> int:
 
     @classmethod
     def _load(cls, path, map_location, **submodules):
-        d = torch.load(path, map_location, weights_only=False)
+        try:
+            d = torch.load(path, map_location, weights_only=False)
+        except AttributeError:
+            logger.error(
+                f"{traceback.format_exc()}\nModel loading failed (full stacktrace above)! It is possible this checkpoint was generated in v2.0 and needs to be converted to v2.1\n Please run 'chemprop convert --conversion v2_0_to_v2_1 -i {path}' and load the converted checkpoint."
+            )
 
         try:
             hparams = d["hyper_parameters"]
diff --git a/chemprop/models/multi.py b/chemprop/models/multi.py
index 012b3fbfd..94d1aae3d 100644
--- a/chemprop/models/multi.py
+++ b/chemprop/models/multi.py
@@ -1,3 +1,4 @@
+import logging
 from typing import Iterable
 
 import torch
@@ -9,6 +10,8 @@
 from chemprop.nn.metrics import ChempropMetric
 from chemprop.nn.transforms import ScaleTransform
 
+logger = logging.getLogger(__name__)
+
 
 class MulticomponentMPNN(MPNN):
     def __init__(
@@ -64,7 +67,12 @@ def get_batch_size(self, batch: MulticomponentTrainingBatch) -> int:
 
     @classmethod
     def _load(cls, path, map_location, **submodules):
-        d = torch.load(path, map_location, weights_only=False)
+        try:
+            d = torch.load(path, map_location, weights_only=False)
+        except AttributeError:
+            logger.error(
+                f"Model loading failed! It is possible this checkpoint was generated in v2.0 and needs to be converted to v2.1\n Please run 'chemprop convert --conversion v2_0_to_v2_1 -i {path}' and load the converted checkpoint."
+            )
 
         try:
             hparams = d["hyper_parameters"]
diff --git a/chemprop/utils/v2_0_to_v2_1.py b/chemprop/utils/v2_0_to_v2_1.py
index 70bee4f96..bc1174d0c 100644
--- a/chemprop/utils/v2_0_to_v2_1.py
+++ b/chemprop/utils/v2_0_to_v2_1.py
@@ -1,3 +1,4 @@
+from os import PathLike
 import pickle
 import sys
 
@@ -33,8 +34,12 @@ def find_class(self, module, name):
         return super().find_class(module, name)
 
 
-if __name__ == "__main__":
+def convert_model_file_v2_0_to_v2_1(model_v1_file: PathLike, model_v2_file: PathLike):
     model = torch.load(
-        sys.argv[1], map_location="cpu", pickle_module=sys.modules[__name__], weights_only=False
+        model_v1_file, map_location="cpu", pickle_module=sys.modules[__name__], weights_only=False
     )
-    torch.save(model, sys.argv[2])
+    torch.save(model, model_v2_file)
+
+
+if __name__ == "__main__":
+    convert_model_file_v2_0_to_v2_1(sys.argv[1], sys.argv[2])
diff --git a/docs/source/tutorial/cli/convert.rst b/docs/source/tutorial/cli/convert.rst
index 12ead9f47..48965eb19 100644
--- a/docs/source/tutorial/cli/convert.rst
+++ b/docs/source/tutorial/cli/convert.rst
@@ -6,5 +6,10 @@ Conversion
 To convert a trained model from Chemprop v1 to v2, run ``chemprop convert`` and specify:
 
  * :code:`--input-path <path>` Path of the Chemprop v1 file to convert.
- * :code:`--output-path <path>` Path where the converted Chemprop v2 will be saved. If unspecified, this will default to ``<CURRENT_DIRECTORY/STEM_OF_INPUT>_v2.ckpt``.
+ * :code:`--output-path <path>` Path where the converted Chemprop v2 model will be saved. If unspecified, this will default to ``<CURRENT_DIRECTORY/STEM_OF_INPUT>_v2.pt``.
+
+To convert a trained model from Chemprop v2.0.x to v2.1.y (or newer), run ``chemprop convert --conversion v2_0_to_v2_1`` and additionally specify:
+
+ * :code:`--input-path <path>` Path of the Chemprop v2.0 file to convert.
+ * :code:`--output-path <path>` Path where the converted Chemprop v2.1 model will be saved. If unspecified, this will default to ``<CURRENT_DIRECTORY/STEM_OF_INPUT>_v2_1.pt``.
 


Reviews:
user: login: KnathanM
body: I think this is useful. I'll try to test it soon.
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-13T21:32:40Z
user: login: KnathanM
body: The conversion works for me. I trained a model in v2.0.3 and tried to predict in the conversion_warning branch. This did not work, but I copied the conversion command provided and the resulting model file did work.

Can you update the docs `docs/source/tutorial/cli/convert.rst`. 

Also see my previous comment about returning the original error as well. 
state: CHANGES_REQUESTED
author_association: MEMBER
submitted_at: 2025-05-07T15:36:42Z
user: login: craabreu
body: I think the notation `v2_0_to_v2_1` can be unclear to `v2.2+` users. Other than that, it looks good to me.
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-08T19:58:47Z
user: login: JacksonBurns
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-15T16:40:12Z
user: login: JacksonBurns
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-15T16:48:02Z
user: login: KnathanM
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-16T20:27:31Z
user: login: KnathanM
state: APPROVED
author_association: MEMBER
submitted_at: 2025-05-18T00:10:21Z

Review comments:
diff_hunk: @@ -253,7 +253,12 @@ def get_batch_size(self, batch: TrainingBatch) -> int:
 
     @classmethod
     def _load(cls, path, map_location, **submodules):
-        d = torch.load(path, map_location, weights_only=False)
+        try:
+            d = torch.load(path, map_location, weights_only=False)
+        except AttributeError:
+            logger.error(
path: chemprop/models/model.py
user: login: KnathanM
body: Might be worth also returning the original error. 
created_at: 2025-03-13T21:29:51Z
author_association: MEMBER
diff_hunk: @@ -5,44 +5,60 @@
 
 from chemprop.cli.utils import Subcommand
 from chemprop.utils.v1_to_v2 import convert_model_file_v1_to_v2
+from chemprop.utils.v2_0_to_v2_1 import convert_model_file_v2_0_to_v2_1
path: chemprop/cli/convert.py
user: login: craabreu
body: Since v2.2 will soon be released, and some users might want to convert models from 2.0 to 2.2, how about generalizing the notation (e.g. `v2_0_to_v2_x` or `v2_0_to_current`)?
created_at: 2025-05-08T19:54:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -253,7 +253,12 @@ def get_batch_size(self, batch: TrainingBatch) -> int:
 
     @classmethod
     def _load(cls, path, map_location, **submodules):
-        d = torch.load(path, map_location, weights_only=False)
+        try:
+            d = torch.load(path, map_location, weights_only=False)
+        except AttributeError:
+            logger.error(
path: chemprop/models/model.py
user: login: JacksonBurns
body: Done in [4bce9f9](https://github.com/chemprop/chemprop/pull/1191/commits/4bce9f9736bf8e61fa4e7c500f018b5cd09bb734)
created_at: 2025-05-15T16:40:12Z
author_association: MEMBER
diff_hunk: @@ -5,44 +5,60 @@
 
 from chemprop.cli.utils import Subcommand
 from chemprop.utils.v1_to_v2 import convert_model_file_v1_to_v2
+from chemprop.utils.v2_0_to_v2_1 import convert_model_file_v2_0_to_v2_1
path: chemprop/cli/convert.py
user: login: JacksonBurns
body: I've thought about which of these would be better, and ultimately I think 2_x and 2_current are about as confusing as 2_1. I will leave the CLI option as is but I have updated the user-facing documentation to clarify that the conversion works for versions of ChemProp including 2.1 _and newer_
created_at: 2025-05-15T16:48:01Z
author_association: MEMBER
diff_hunk: @@ -6,5 +6,10 @@ Conversion
 To convert a trained model from Chemprop v1 to v2, run ``chemprop convert`` and specify:
 
  * :code:`--input-path <path>` Path of the Chemprop v1 file to convert.
- * :code:`--output-path <path>` Path where the converted Chemprop v2 will be saved. If unspecified, this will default to ``<CURRENT_DIRECTORY/STEM_OF_INPUT>_v2.ckpt``.
+ * :code:`--output-path <path>` Path where the converted Chemprop v2 will be saved. If unspecified, this will default to ``<CURRENT_DIRECTORY/STEM_OF_INPUT>_v2.pt``.
+
+To convert a trained model from Chemprop v2.0.x to v2.1.y (or newer), run ``chemprop convert --conversion v2_0_to_v2_1`` and additionally specify:
+
+ * :code:`--input-path <path>` Path of the Chemprop v1 file to convert.
path: docs/source/tutorial/cli/convert.rst
user: login: KnathanM
body: ```suggestion
 * :code:`--input-path <path>` Path of the Chemprop v2.0 file to convert.
```
created_at: 2025-05-16T20:22:15Z
author_association: MEMBER
diff_hunk: @@ -6,5 +6,10 @@ Conversion
 To convert a trained model from Chemprop v1 to v2, run ``chemprop convert`` and specify:
 
  * :code:`--input-path <path>` Path of the Chemprop v1 file to convert.
- * :code:`--output-path <path>` Path where the converted Chemprop v2 will be saved. If unspecified, this will default to ``<CURRENT_DIRECTORY/STEM_OF_INPUT>_v2.ckpt``.
+ * :code:`--output-path <path>` Path where the converted Chemprop v2 will be saved. If unspecified, this will default to ``<CURRENT_DIRECTORY/STEM_OF_INPUT>_v2.pt``.
+
+To convert a trained model from Chemprop v2.0.x to v2.1.y (or newer), run ``chemprop convert --conversion v2_0_to_v2_1`` and additionally specify:
+
+ * :code:`--input-path <path>` Path of the Chemprop v1 file to convert.
+ * :code:`--output-path <path>` Path where the converted Chemprop v2 will be saved. If unspecified, this will default to ``<CURRENT_DIRECTORY/STEM_OF_INPUT>_v2_1.pt``.
path: docs/source/tutorial/cli/convert.rst
user: login: KnathanM
body: ```suggestion
 * :code:`--output-path <path>` Path where the converted Chemprop v2.1 model will be saved. If unspecified, this will default to ``<CURRENT_DIRECTORY/STEM_OF_INPUT>_v2_1.pt``.
```
created_at: 2025-05-16T20:22:58Z
author_association: MEMBER
diff_hunk: @@ -6,5 +6,10 @@ Conversion
 To convert a trained model from Chemprop v1 to v2, run ``chemprop convert`` and specify:
 
  * :code:`--input-path <path>` Path of the Chemprop v1 file to convert.
- * :code:`--output-path <path>` Path where the converted Chemprop v2 will be saved. If unspecified, this will default to ``<CURRENT_DIRECTORY/STEM_OF_INPUT>_v2.ckpt``.
+ * :code:`--output-path <path>` Path where the converted Chemprop v2 will be saved. If unspecified, this will default to ``<CURRENT_DIRECTORY/STEM_OF_INPUT>_v2.pt``.
path: docs/source/tutorial/cli/convert.rst
user: login: KnathanM
body: ```suggestion
 * :code:`--output-path <path>` Path where the converted Chemprop v2 model will be saved. If unspecified, this will default to ``<CURRENT_DIRECTORY/STEM_OF_INPUT>_v2.pt``.
```
created_at: 2025-05-16T20:23:20Z
author_association: MEMBER

PR comments:
user: login: JacksonBurns
created_at: 2025-05-06T20:57:49Z
body: TODO: add mention in the description for the args that one must use the v1 featurizer with v1 to v2 converted models
author_association: MEMBER
user: login: JacksonBurns
created_at: 2025-05-15T16:49:16Z
body: @craabreu @KnathanM I have addessed all of your comments to the best of my ability, please re-review and let me know how everything looks!
author_association: MEMBER