From d13629cb26cbd9398ad79ac0656970cc64e2c6af Mon Sep 17 00:00:00 2001
From: Jacky Lee <39754370+jla524@users.noreply.github.com>
Date: Wed, 10 May 2023 09:01:22 -0700
Subject: [PATCH] ResNet: match implementation with Nvidia and PyTorch (#770)

* Match ResNet implementation with pytorch and nvidia

* Reduce number of Epochs
---
 examples/train_resnet.py | 2 +-
 models/resnet.py         | 6 +++---
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/examples/train_resnet.py b/examples/train_resnet.py
index c7b2f7a8a1..970eeba333 100755
--- a/examples/train_resnet.py
+++ b/examples/train_resnet.py
@@ -36,7 +36,7 @@ if __name__ == "__main__":
     lambda x: x / 255.0,
     lambda x: np.tile(np.expand_dims(x, 1), (1, 3, 1, 1)).astype(np.float32),
   ])
-  for _ in range(10):
+  for _ in range(5):
     optimizer = optim.SGD(optim.get_parameters(model), lr=lr, momentum=0.9)
     train(model, X_train, Y_train, optimizer, 100, BS=32, transform=transform)
     evaluate(model, X_test, Y_test, num_classes=classes, transform=transform)
diff --git a/models/resnet.py b/models/resnet.py
index f4c498c982..0455331a64 100644
--- a/models/resnet.py
+++ b/models/resnet.py
@@ -27,6 +27,7 @@ class BasicBlock:
 
 
 class Bottleneck:
+  # NOTE: the original implementation places stride at the first convolution (self.conv1), this is the v1.5 variant
   expansion = 4
 
   def __init__(self, in_planes, planes, stride=1):
@@ -34,7 +35,7 @@ class Bottleneck:
     self.bn1 = nn.BatchNorm2d(planes)
     self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1, stride=stride, bias=False)
     self.bn2 = nn.BatchNorm2d(planes)
-    self.conv3 = nn.Conv2d(planes, self.expansion *planes, kernel_size=1, bias=False)
+    self.conv3 = nn.Conv2d(planes, self.expansion*planes, kernel_size=1, bias=False)
     self.bn3 = nn.BatchNorm2d(self.expansion*planes)
     self.downsample = []
     if stride != 1 or in_planes != self.expansion*planes:
@@ -52,7 +53,6 @@ class Bottleneck:
     return out
 
 class ResNet:
-  # def __init__(self, block, num_blocks, num_classes=10, url=None):
   def __init__(self, num, num_classes):
     self.num = num
 
@@ -76,7 +76,7 @@ class ResNet:
 
     self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, bias=False, padding=3)
     self.bn1 = nn.BatchNorm2d(64)
-    self.layer1 = self._make_layer(self.block, 64, self.num_blocks[0], stride=2)
+    self.layer1 = self._make_layer(self.block, 64, self.num_blocks[0], stride=1)
     self.layer2 = self._make_layer(self.block, 128, self.num_blocks[1], stride=2)
     self.layer3 = self._make_layer(self.block, 256, self.num_blocks[2], stride=2)
     self.layer4 = self._make_layer(self.block, 512, self.num_blocks[3], stride=2)