PyTorch: Load/unpickle only the state dict of a model saved with torch.save

user_5128131 0
Viewed 345

I saved a nn.Module model using (logically):

model = MyWeirdModel()
model.patched_features = .....
train(model)
torch.save(model, file)

Ideally, one would load this model using

model = torch.load(file)

However, in my case this doesn't work because Pickle uses the static class definition when un-pickling, so I get AttributeError: 'MyWeirdModel' object has no attribute 'patched_features' (this attribute was added to MyWeirdModel at runtime).

I would like to avoid having to re-train the model, so I don't want to change the code for saving, only loading.

# Initialise the model in the same way as before
model = MyWeirdModel()
model.patched_features = .....

state_dict = load_state_dict_only(file) # How does one do this?

model.load_state_dict(state_dict)

My understanding is that torch.save() saves the model AND the state dict. How do I load only the state dict from the pickled model, such that I can recover the model?

pytorch pickle python
1 Answers

Referring to the documentation,

only layers with learnable parameters (convolutional layers, linear layers, etc.) and registered buffers (batchnorm’s running_mean) have entries in the model’s state_dict.

Your patched_features is not stored in the model’s state_dict, but you can do so by registering it with register_parameter as pointed out by ptrblck in this post but it may involve retraining.

user_101915540
Related