Add save function to GNNModel

CHANGELOG.md

@@ -25,6 +25,7 @@ The rules for this file:
 - @mattwthompson
 
 ### Added
+- `GNNModel.save` function (PR #29)
 - `GNNModel.load` function (PR #26)
 - `convolution_dropout` and `readout_dropout` keywords to GNNModel (PR #26)
 

openff/nagl/nn/_models.py

@@ -13,6 +13,9 @@
     from openff.nagl.features.bonds import BondFeature
     from openff.nagl.molecule._dgl.batch import DGLMoleculeBatch
     from openff.nagl.molecule._dgl.molecule import DGLMolecule
+    from openff.nagl.nn.postprocess import PostprocessLayerMeta
+    from openff.nagl.nn.activation import ActivationFunction
+    from openff.nagl.nn.gcn._base import GCNStackMeta
 
 
 def rmse_loss(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
@@ -113,6 +116,52 @@ def _as_nagl(self):
 
 
 class GNNModel(BaseGNNModel):
+    """
+    A model that applies a graph convolutional step followed by
+    pooling and readout steps.
+
+    Parameters
+    ----------
+    convolution_architecture: Union[str, GCNStackMeta]
+        The graph convolution architecture.
+        This can be given either as a class,
+        e.g. :class:`~openff.nagl.nn.gcn._sage.SAGEConvStack`
+        or as a string, e.g. ``"SAGEConv"``.
+    n_convolution_hidden_features: int
+        The number of features in each of the hidden convolutional layers.
+    n_convolution_layers: int
+        The number of hidden convolutional layers to generate. These are the 
+        layers in the convolutional module between the input layer and the 
+        pooling layer.
+    n_readout_hidden_features: int
+        The number of features in each of the hidden readout layers.
+    n_readout_layers: int
+        The number of hidden readout layers to generate. These are the layers
+        between the convolution module's pooling layer and the readout module's
+        output layer. The pooling layer may be considered to be both the
+        convolution module's output layer and the readout module's input layer.
+    activation_function: Union[str, ActivationFunction]
+        The activation function to use for the readout module.
+        This can be given either as a class,
+        e.g. :class:`~openff.nagl.nn.activation.ActivationFunction.ReLU`,
+        or as a string, e.g. ``"ReLU"``.
+    postprocess_layer: Union[str, PostprocessLayerMeta]
+        The postprocess layer to use.
+        This can be given either as a class,
+        e.g. :class:`~openff.nagl.nn.postprocess.ComputePartialCharges`,
+        or as a string, e.g. ``"compute_partial_charges"``.
+    atom_features: Tuple[AtomFeature, ...]
+        The atom features to use.
+    bond_features: Tuple[BondFeature, ...]
+        The bond features to use.
+    loss_function: Callable[[torch.Tensor, torch.Tensor], torch.Tensor]
+        The loss function. This is RMSE by default, but can be any function
+        that takes a predicted and target tensor and returns a scalar loss tensor.
+    convolution_dropout: float
+        The dropout probability to use in the convolutional layers.
+    readout_dropout: float
+        The dropout probability to use in the readout layers.
+    """
     @classmethod
     def from_yaml_file(cls, *paths, **kwargs):
         import yaml
@@ -134,13 +183,13 @@ def n_atom_features(self) -> int:
 
     def __init__(
         self,
-        convolution_architecture: str,
+        convolution_architecture: Union[str, "GCNStackMeta"],
         n_convolution_hidden_features: int,
         n_convolution_layers: int,
         n_readout_hidden_features: int,
         n_readout_layers: int,
-        activation_function: str,
-        postprocess_layer: str,
+        activation_function: Union[str, "ActivationFunction"],
+        postprocess_layer: Union[str, "PostprocessLayerMeta"],
         readout_name: str,
         learning_rate: float,
         atom_features: Tuple["AtomFeature", ...],
@@ -201,6 +250,21 @@ def __init__(
     def compute_property(
         self, molecule: "Molecule", as_numpy: bool = False
     ) -> "torch.Tensor":
+        """
+        Compute the trained property for a molecule.
+
+        Parameters
+        ----------
+        molecule: :class:`~openff.toolkit.topology.Molecule`
+            The molecule to compute the property for.
+        as_numpy: bool
+            Whether to return the result as a numpy array.
+            If ``False``, the result will be a ``torch.Tensor``.
+
+        Returns
+        -------
+        result: torch.Tensor or numpy.ndarray
+        """
         try:
             values = self._compute_property_dgl(molecule)
         except MissingOptionalDependencyError:
@@ -260,9 +324,37 @@ def _validate_features(features, feature_class):
 
     @classmethod
     def load(cls, model: str, eval_mode: bool = True):
-        import torch
-
-        model_kwargs = torch.load(model)
+        """
+        Load a model from a file.
+
+        Parameters
+        ----------
+        model: str
+            The path to the model to load.
+            This should be a file containing a dictionary of
+            hyperparameters and a state dictionary,
+            with the keys "hyperparameters" and "state_dict".
+            This can be created using the `save` method.
+        eval_mode: bool
+            Whether to set the model to evaluation mode.
+
+        Returns
+        -------
+        model: GNNModel
+
+        Examples
+        --------
+
+        >>> model.save("model.pt")
+        >>> new_model = GNNModel.load("model.pt")
+
+        Notes   
+        -----
+        This method is not compatible with normal Pytorch
+        models saved with ``torch.save``, as it expects
+        a dictionary of hyperparameters and a state dictionary.
+        """
+        model_kwargs = torch.load(str(model))
         if isinstance(model_kwargs, dict):
             model = cls(**model_kwargs["hyperparameters"])
             model.load_state_dict(model_kwargs["state_dict"])
@@ -274,3 +366,31 @@ def load(cls, model: str, eval_mode: bool = True):
             model.eval()
 
         return model
+
+    def save(self, path: str):
+        """
+        Save this model to a file.
+
+        Parameters
+        ----------
+        path: str
+            The path to save this file to.
+
+        Examples
+        --------
+
+        >>> model.save("model.pt")
+        >>> new_model = GNNModel.load("model.pt")
+
+        Notes
+        -----
+        This method writes a dictionary of the hyperparameters and the state dictionary,
+        with the keys "hyperparameters" and "state_dict".
+        """
+        torch.save(
+            {
+                "hyperparameters": self.hparams,
+                "state_dict": self.state_dict(),
+            },
+            str(path),
+        )

openff/nagl/tests/nn/test_model.py

@@ -194,3 +194,11 @@ def test_load(self, openff_methane_uncharged):
         charges = model.compute_property(openff_methane_uncharged, as_numpy=True)
         expected = np.array([-0.111393,  0.027848,  0.027848,  0.027848,  0.027848])
         assert_allclose(charges, expected, atol=1e-5)
+
+    def test_save(self, am1bcc_model, openff_methane_uncharged, tmpdir):
+        with tmpdir.as_cwd():
+            am1bcc_model.save("model.pt")
+            model = GNNModel.load("model.pt", eval_mode=True)
+            charges = model.compute_property(openff_methane_uncharged, as_numpy=True)
+            expected = np.array([-0.143774, 0.035943, 0.035943, 0.035943, 0.035943])
+            assert_allclose(charges, expected, atol=1e-5)