Fixed the input image size for the resnet18

This commit is contained in:
Yunjey Choi
2018-02-15 22:48:51 +09:00
committed by GitHub
parent 655c6df3ba
commit f73f67e310

View File

@ -150,7 +150,7 @@ for param in resnet.parameters():
resnet.fc = nn.Linear(resnet.fc.in_features, 100) # 100 is for example.
# For test.
images = Variable(torch.randn(10, 3, 256, 256))
images = Variable(torch.randn(10, 3, 224, 224))
outputs = resnet(images)
print (outputs.size()) # (10, 100)
@ -162,4 +162,4 @@ model = torch.load('model.pkl')
# Save and load only the model parameters(recommended).
torch.save(resnet.state_dict(), 'params.pkl')
resnet.load_state_dict(torch.load('params.pkl'))
resnet.load_state_dict(torch.load('params.pkl'))