预训练模型的预加载
预训练模型的预加载
提问于 2021-05-02 14:44:57
我用以下代码保存了我的模型：
from google.colab import files
torch.save(net, 'model.pth')
# download checkpoint file
files.download('model.pth')然后以这种方式上传并检查图像(x)：
model = torch.load('model.pth')
model.eval()
torch.argmax(model(x))在以前的会话中，它工作得很好，但是后来我启动了一个新会话，并尝试了上面的代码，并得到了这样一个错误：
AttributeError: Can't get attribute 'EfficientNet' on <module '__main__'>也许有人知道怎么处理？
回答 2
Data Science用户
回答已采纳
发布于 2021-05-02 16:11:45
最近也犯了同样的错误。
您的网络通常被定义为一个类(此处为class EfficientNet(nn.Module)。似乎当我们加载一个模型时，它需要定义类以便它能够实例化它。
在我的例子中，类是在培训.py文件中定义的。因此，我所做的修复该错误的只是复制粘贴(导入它似乎对我无效，所以我不得不将整个类定义复制到加载模型的文件中)。
在我的示例中，类定义类似于此，需要在使用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')类基本上定义了模型的体系结构，如果没有定义，那么load函数就不知道在哪里放置权重。我的是用torch.load(path)而不是model.load_state_dict(torch.load(path))做的，但这不应该有什么区别。
Data Science用户
发布于 2021-05-02 16:20:41
您需要首先定义模型。一旦定义了模型，那么，实例化其中的一个类。完成之后，使用model.load_state_dict(torch.load(path_to_model_file))。
页面原文内容由Data Science提供。腾讯云小微IT领域专用引擎提供翻译支持
原文链接：
https://datascience.stackexchange.com/questions/93863
复制相关文章
相似问题
腾讯云开发者
