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