I have trained a model and save model using torch.save. Then after training I have loaded the model using train.load but I am getting this error
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'
Here is my code for test part
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)
I have commented training part of the code in the file, so in a way this and loading the data part is all that is there in the file now.
What am I doing wrong?
Like @jodag has said. you probably have saved a state_dict instead of a model, which is recommended by the community as well.
This link explains the difference between two. To keep my answer self contained, I copy the snippet from the documentation. Here is the recommended way:
Save:
torch.save(model.state_dict(), PATH)
Load:
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()
You could also save the entire model instead of saving the state_dict, if you really need to use the model the way you do.
Save:
torch.save(model, PATH)
Load:
# Model class must be defined somewhere
model = torch.load(PATH)
model.eval()
The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.