algorithm/pytorch-tutorial
1
0
Fork 0
mirror of https://github.com/yunjey/pytorch-tutorial.git synced 2026-03-13 09:11:37 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #141 from keineahnung2345/fix-resnet

Browse Source 
[Bug fix] redundant layers in ResNet
This commit is contained in:
Yunjey Choi
2018-11-06 22:04:38 +09:00
committed by GitHub
parent 6f28678302 606d0aa188
commit c48008bf9d
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
tutorials/02-intermediate/deep_residual_network/main.py
View File

@@ -82,8 +82,8 @@ class ResNet(nn.Module):
self.bn = nn.BatchNorm2d(16)
self.relu = nn.ReLU(inplace=True)
self.layer1 = self.make_layer(block, 16, layers[0])
self.layer2 = self.make_layer(block, 32, layers[0], 2)
self.layer3 = self.make_layer(block, 64, layers[1], 2)
self.layer2 = self.make_layer(block, 32, layers[1], 2)
self.layer3 = self.make_layer(block, 64, layers[2], 2)
self.avg_pool = nn.AvgPool2d(8)
self.fc = nn.Linear(64, num_classes)
@@ -112,7 +112,7 @@ class ResNet(nn.Module):
out = self.fc(out)
return out
model = ResNet(ResidualBlock, [2, 2, 2, 2]).to(device)
model = ResNet(ResidualBlock, [2, 2, 2]).to(device)
# Loss and optimizer
@@ -166,4 +166,4 @@ with torch.no_grad():
print('Accuracy of the model on the test images: {} %'.format(100 * correct / total))
# Save the model checkpoint
torch.save(model.state_dict(), 'resnet.ckpt')
torch.save(model.state_dict(), 'resnet.ckpt')