Load parameters only from an entirely saved model

vision
1

Hi, I have trained a neural network for Purpose 1 and I mistakenly use torch.save(model, PATH), which saves the entire model rather than only weights.

As the performance is very good, I would like to plug this model into another model to assist Purpose 2…

I have created the exact model object…
Then, I try to load the parameters by

checkpoint = torch.load("./models/xxxnet_weights.pk", map_location=torch.device('cpu'))
model.load_state_dict(checkpoint['model_state_dict'])

But, the torch.load gives an error like this

AttributeError: Can't get attribute 'xxxNet' on <module 'models' from '/home/xxx/Projects/xxxx/models/__init__.py'>

(As stated in the official tutorial https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict) I suspect this is due to the torch.save saves the directory information, which is different now…

Is there a way that I can still load these parameters?

Many thanks

2

I think I can provide a solution for myself.

Load the model in the original project folder, save using torch.save(model.state_dict(), PATH) with parameters only again to another checkpoint. Then, load this new one in the new project.