Error loading saved model


(Shiv) #1

I have a model:

cnn=CNN()

torch.save(cnn, ‘./model/trained_model.pt’)

model = torch.load(’./model/trained_model.pt’)

AttributeError: Can’t get attribute ‘CNN’ on <module ‘main’>

why am I getting this error and how to solve it?


(Russell Wong) #2

I was having the similar error before.
How I solved it by copying the code of CNN class to the same file that load the model.
I think some hows the namespace is hard code when you save the model.
So then it can’t search when you load the model again.


(Shiv) #3

Thanks @russellwmy. It worked :slight_smile:


(Russell Wong) #4

Good, But the way to solve it is … so bad :sweat_smile:


(Nasib Ullah) #5

Thank you @russellwmy .Anyone found any other way?


(Thong Nguyen) #6

Hi, I also have this problem and figure it out and just wanna post my solution here in case it helps anyone.

I read somewhere that Pytorch use Pickle to take care of the save/load model, the problem is really about Pickle, not Pytorch itself.

The solution is really simple, you just have to explicitly import the class definition. It’s equivalent to copy and paste the class definition.

In my case I have the following directory structure:

Release-code
|----- myclass.py
Application
|----- app1.py

In app1.py, I make sure ./Release-code/ is in python path, and I just have to say

from myclass import class_definition

then torch.load(model_name) works

If your class_definition calls other classes, you also need to import them too.

This thread explains the problem and solutions for it: https://stackoverflow.com/questions/27732354/unable-to-load-files-using-pickle-and-multipile-modules