PyTorch: saving both weights and model definition
During prototyping, I often perform numerous changes to a PyTorch model. For instance, suppose the first model I am experimenting with is:
class Model(nn.Module):
    def __init__(self, **args):
        super().__init__()
        self.l1 = nn.Linear(128, 1)

then I will add another layer:
class Model(nn.Module):
    def __init__(self, **args):
        super().__init__()
        self.l1 = nn.Linear(128, 32)
        self.l2 = nn.Linear(32, 1)

or maybe add some convolutions, and so on.
The problem is that I often get disorganized the more experiments I perform, as I haven't found a straightforward way of saving both model definition and its weights so that I can load a previous state.
I know I can do:
torch.save({'model': Model(), 'state': model.state_dict()}, path)
# or directly
torch.save(model, path)

but then loading the model also requires that the model class (here, Model) to exist in the current file.
In Keras you can simply do:
model = ...  # Get model (Sequential, Functional Model, or Model subclass)
model.save('path/to/location')

which saves model's architecture/config and weights, among other things. This means that you can load the model without having defined the architecture:
model = keras.models.load_model('path/to/location')

Refering to Keras model saving:

The SavedModel and HDF5 file contains:

the model's configuration (topology)
the model's weights
the model's optimizer's state (if any)

Thus models can be reinstantiated in the exact same state, without any
of the code used for model definition or training.

This is what I want to achieve in PyTorch.
Is there a similar approach for PyTorch? What is the best practice for these situations?

As Pytorch provides a huge amount of flexibility in the model, it will be challenging to save the architecture along with the weights in a single file. Keras models are usually built solely by stacking keras components, but pytorch models are orchestrated by the library consumer in their own way and therefore can contain any sort of logic.
I think you have three choices:

Come up with a organised schema for your experiments so that losing the model definition is less likely. You could go for something as simple as a file named through a schema that only defines each model. I would recommend this approach as this level of organisation would likely benefit your prototyping in other ways and the overhead is minimal.

Try and save the code along with the pickle file. Although potentially possible, I think this would lead you down a rabbit-hole with a lot of potential problems.

Use a different standardised way of saving the model, such as onnx. I would recommend this route if you do not want to go with option 1. Onnx does allow you to save a pytorch model's architecture along with its weights but comes with a few drawbacks. For example, it only supports some operations so completely custom forward methods or use of non-matrix operations may not work.


