我已经训练了一个模型,并使用torch.save保存模型。然后在训练之后,我使用train.load加载了模型,但是我得到了这个错误
Traceback (most recent call last):
File "/home/fsdfs.py", line 219, in <module>
test(model, 'cuda', testloader)
File "/home/fsdfs.py", line 201, in test
model.eval()
AttributeError: 'collections.OrderedDict' object has no attribute 'eval'
字符串
下面是我的测试部分代码
model = torch.load("train_5.pth")
def test(model, device, test_loader):
model.eval()
test_loss = 0
correct = 0
with torch.no_grad():
for data, target in test_loader:
data, target = data.to('cuda'), target.to('cuda')
output = model(data)
#test_loss += f.cross_entropy(output, target, reduction='sum').item() # sum up batch loss
pred = output.argmax(1, keepdim=True) # get the index of the max log-probability
print(pred, target)
correct += pred.eq(target.view_as(pred)).sum().item()
test_loss /= len(test_loader.dataset)
print('\nTest set: Accuracy: {}/{} ({:.0f}%)\n'.format(
correct, len(test_loader.dataset),
100. * correct / len(test_loader.dataset)))
test(model, 'cuda', testloader)
型
我已经在文件中注解了代码的训练部分,所以在某种程度上,这部分和加载数据部分就是文件中的全部内容。
我做错了什么?
1条答案
按热度按时间rekjcdws1#
就像@jodag说的。你可能保存了一个state_dict而不是一个模型,这也是社区推荐的。
这个链接解释了两者之间的区别。为了保持我的答案是独立的,我从文档中复制了这个片段。以下是推荐的方法:
保存:
字符串
加载:
型
如果你真的需要这样使用模型,你也可以保存整个模型,而不是保存state_dict。
保存:
型
加载:
型