How can I load a model in pytorch without having to remember the parameters used?

Question

I am training a model in pytorch for which I have made a class like so:

from torch import nn

class myNN(nn.Module):
    def __init__(self, dense1=128, dense2=64, dense3=32, ...):
         self.MLP = nn.Sequential(
            nn.Linear(dense1, dense2),
            nn.ReLU(),
            nn.Linear(dense2, dense3),
            nn.ReLU(),
            nn.Linear(dense3, 1)
         )
         ...

In order to save it I am using:

torch.save(model.state_dict(), checkpoint_model_path)

and to load it I am using:

model = myNN()   # or with specified parameters
model.load_state_dict(torch.load(model_file))

However, in order for this method to work I have to use the right values in myNN()'s constructor. That means that I would need to somehow remember or store which parameters (layer sizes) I have used in each case in order to properly load different models.

Is there a flexible way to save/load models in pytorch where I would also read the size of the layers?

E.g. by loading a myNN() object directly or somehow reading the layer sizes from the saved pickle file?

I am hesitant to try the second method in Best way to save a trained model in PyTorch? due to the warnings mentioned there. Is there a better way to achieve what I want?

score 11 · Accepted Answer · answered Nov 09 '21 at 21:23

Indeed serializing the whole Python is quite a drastic move. Instead, you can always add user-defined items in the saved file: you can save the model's state along with its class parameters. Something like this would work:

First save your arguments in the instance such that we can serialize them when saving the model:

class myNN(nn.Module):
    def __init__(self, dense1=128, dense2=64, dense3=32):
        super().__init__()
        self.kwargs = {'dense1': dense1, 'dense2': dense2, 'dense3': dense3}
        self.MLP = nn.Sequential(
            nn.Linear(dense1, dense2),
            nn.ReLU(),
            nn.Linear(dense2, dense3),
            nn.ReLU(),
            nn.Linear(dense3, 1))

We can save the parameters of the model along with its initializer arguments:
```
>>> torch.save([model.kwargs, model.state_dict()], path)
```

Then load it:

>>> kwargs, state = torch.load(path)
>>> model = myNN(**kwargs)
>>> model.load_state_dict(state)
<All keys matched successfully>

Works like a charm! Thank you, now I don't have to remember what I used for each saved model! — Michael, Nov 11 '21 at 10:05

How can I load a model in pytorch without having to remember the parameters used?

1 Answers1