Summary

Currently some small tensors saved twice in a chemprop model file, both in hyper_parameters and in state_dict. I think there isn't a clean way to not double save them and the small tensors are small enough to not worry about, so the best thing to do is to continue double saving, suppress any warnings, and be explicit that is what we are doing.

The warning can be suppressed by adding the small nn.Modules arguments to self.save_hyperparameters(ignore=...) and then manually adding them to hparams. #832 started this for criterion and output_transform of _FFNPredictorBase.

The warning looks something like "Attribute '...' is an instance of nn.Module and is already saved during checkpointing. It is recommended to ignore them using self.save_hyperparameters(ignore=['...'])"

Background:

All of this is related to the question of "How do we reload a previous model?". Closely related is the question "How do we save a model so that it can be reloaded?". There are two parts of models that have to be saved and loaded, the architecture and the "numbers" (weights, biases, etc.).

MPNN is a lightning.pytorch.LightningModule which inherits from both lightning.pytorch.core.mixins.HyperparametersMixin and torch.nn.Module. HyperparametersMixin takes care of saving the model architecture (by saving the input arguments when creating the model) and Module takes care of saving the "numbers".

I'll discuss three representative attributes of MPNN that are each saved and loaded differently.

message_passing

architecture

HyperparametersMixin has a method save_hyperparameters which will add all the arguments passed to __init__ to self.hparams. In the case of message_passing we already include it in the self.save_hyperparameters(ignore=...) ignore list because message_passing is an object that contains all the weights that are already saved in the state_dict. save_hyperparameters would just pickle the whole object leading to the weights being saved in both hparam and in state_dict.

We want to just save the instructions for how to make a new message_passing module with the same architecture as the saved version. It will be created with random weights, but then those will be overwritten with the correct weights when the state_dict is loaded. To do this _MessagePassingBase also inherits from HyperparametersMixin and saves its own hparams. This hparams dict is then used as the value for the message_passing key in MPNN's hparams.

weights

Module has an overloaded __setattr__ which affects the behavior setting attributes. When we do self.message_passing = message_passing, the message passing object gets registered as a submodule of MPNN (because it is also a Module) and the state_dict for MPNN will include the parameters of message_passing. (I think that is done here.) message_passing will also appear in the output of print(model).

To then load an MPNN, the message_passing will first be created with random weights using the hparams and then the state_dict is loaded which overwrites the message_passing weights with the correct ones.

metrics

The metrics given to MPNN is not a Module. It is a list whose elements are Metrics which are Modules. This means that self.metrics = metrics does not register metrics as a submodule of MPNN so the numbers in the individual Metric state_dicts are not added to MPNNs state_dict. This also means that metrics does not appear in print(model) which seems reasonable to me. metrics is not included in the ignore list of save_hyperparameters so the whole thing is pickled and reloaded when saving and loading a model.

X_d_transform

We want X_d_transform to be a submodule of MPNN so that it appears in print(model) among other things. We do this by setting it as an attribute (that also includes a check if it is None) self.X_d_transform = X_d_transform if X_d_transform is not None else nn.Identity(). If X_d_transformis notNone, it's state_dict(which includes the scale and mean) will be added toMPNNs state_dict`.

The question is, do we include the whole X_d_transform object in hparams of MPNN or not? I think we must because load_from_file of MPNN doesn't know to initialize a new X_d_transform like it did with message_passing. But if we do include the whole object, the scale and mean of X_d_transform will be duplicated in both hparams and in state_dict.

Our options are

1. Duplicate and suppress warnings. I favor this because the duplicated arrays are small.
2. Make _ScaleTransformMixin more like _MessagePassingBase by having it inherit from HyperparametersMixin and save it's input arguments in hparams. This seems more complicated to me.
3. Exclude scale and mean from the state_dict by not registering them as buffers. I remember that it was a concious decision to include them in the state_dict so this doesn't seem like a good solution.

These same options and questions apply to the other inputs that get this warning from lightning including: "V_d_transform", "graph_transform", "criterion", and "output_transform"

Other note

One potential issue with #832 that I changed is _FFNPredictorBase expects output_transform to be an UnscaleTransform or None, but output_transform is saved after converting Nones to nn.Identity. If we ever added some additional logic to check if output_transform is input as None, this would break that.