0
$\begingroup$

I saved my model with this code:

from google.colab import files
torch.save(net, 'model.pth')

# download checkpoint file
files.download('model.pth')

Then uploaded this way and checked on an image (x):

model = torch.load('model.pth')
model.eval()
torch.argmax(model(x))

And on the old session, it worked great, but then I started a new session and tried the code above, and got such an error:

AttributeError: Can't get attribute 'EfficientNet' on <module '__main__'>

Maybe somebody knows how to deal with it?

Improve this question
$\endgroup$
1
$\begingroup$

Got the very same error recently.

Your network is usually defined as a class (here class EfficientNet(nn.Module). It seems when we load a model, it needs the class to be defined so it can instantiate it.

In my case, the class was defined in the training .py file. So what I did to fix that error was just copy-paste (it seems importing it didn't work for me, so I had to copy paste) the whole class definition to the file where you load your model.

In my case the class definition looked something like that and needed to be pasted before using the torch.load function:

class FCNClass(nn.Module):
    def __init__(self, in_feat=1000, nb_classes=3, nb_hid=1, hidden_size=500, act=nn.ReLU):
        super(FCNClass, self).__init__()

        self.act = act()
        self.flat = nn.Flatten()

        self.fc1 = nn.Linear(in_feat, hidden_size)
        self.fcs = nn.ModuleList()
        for nb in range(nb_hid):
            self.fcs.add_module("hid" + str(nb), nn.Linear(hidden_size, hidden_size))
        self.out = nn.Linear(hidden_size, nb_classes)

    def forward(self, x):
        x = self.flat(x)
        x = self.fc1(x)
        x = self.act(x)
        for lay in self.fcs:
            x = lay(x)
            x = self.act(x)
        x = self.out(x)
        x = self.act(x)
        return x


torch.load('Model.h5')

The class basically defines the architecture of the model, if it is not defined, then the load function doesn't know where to put the weights. Mine was done with torch.load(path) and not model.load_state_dict(torch.load(path)) but it should not make a difference

Improve this answer
$\endgroup$
3
  • $\begingroup$ This command "model.load_state_dict(torch.load(path_to_model_file))" loading only pre-trained weights as I understand, but my cnn class should be described where I uploading it? Please confirm and thanks for your answer $\endgroup$ – Adolf Miszka May 2 at 16:33
  • 1
    $\begingroup$ post edited, hopes it clarify your understanding, pytorch.org/tutorials/beginner/…. Here is the tutorial for loading models with state dictionnary. Maybe the issue in your code is when you save your network, be sure to save the state_dict if you load the state_dict : torch.save(model.state_dict(), PATH) $\endgroup$ – Ubikuity May 2 at 16:51
  • $\begingroup$ Nope, the issue was like yours :) I was thinking the whole model is saved not just weights, anyway thanks. $\endgroup$ – Adolf Miszka May 2 at 17:01
1
$\begingroup$

You need to first define the model. Once you have defined the model, then, instantiate a class of it. Once that is done, use model.load_state_dict(torch.load(path_to_model_file)).

Improve this answer
$\endgroup$
1
  • $\begingroup$ This command "model.load_state_dict(torch.load(path_to_model_file))" loading only pre-trained weights as I understand, but my cnn class should be described where I uploading it? Please confirm and thanks for your answer @Abhishek Verma $\endgroup$ – Adolf Miszka May 2 at 16:29

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.