How to load an older model in PyTorch 1.6?

deployment
1

I have some models trained in PyTorch 1.5 and saved by torch.save(filename, model), instead of torch.save(filename, model.state_dict()).

Now, I am not able to load them in PyTorch 1.6.
I get torch.nn.modules.module.ModuleAttributeError: 'BatchNorm2d' object has no attribute '_non_persistent_buffers_set'

What is the correct way to load such models?

Following https://github.com/ultralytics/yolov5/issues/58, I tried this on the model and seemed to load without problem. But is this the right way?

model = ModelDenseNet()
model1 = torch.load(model_file)
for k, m in model1.named_modules():
   m._non_persistent_buffers_set = set()  # pytorch 1.6.0 compatability
model.load_state_dict(model1.state_dict())

Thanks

2

I would consider this a workaround and would try to make sure your model is still working as intended.
E.g. you could create some tests using predefined data and compare the outputs, gradients, losses etc.
The recommended way would be to store the state_dict instead of the model directly.

3

The workaround seems to work (gives the same accuracy on the test set). I was wondering if there is a better way.

Would be great if PyTorch provided and option to torch.save/torch.load to save load either the whole model or the state_dict.

4

The recommended way of storing the state_dict should be working fine or are you seeing any issues with it? As explained in the Serialization semantics, storing and loading the model directly could be breaking in various ways due to the way Python pickles the model.