I have some models trained in PyTorch 1.5 and saved by
torch.save(filename, model), instead of
Now, I am not able to load them in PyTorch 1.6.
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
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.
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.
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.