Simplify resnet definition

Author: ruotianluo

lib/nets/resnet_v1.py

@@ -19,137 +19,23 @@
 import math
 import torch.utils.model_zoo as model_zoo
 
+import torchvision
+from torchvision.models.resnet import BasicBlock, Bottleneck
 
-__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
-       'resnet152']
-
-
-model_urls = {
-  'resnet18': 'https://s3.amazonaws.com/pytorch/models/resnet18-5c106cde.pth',
-  'resnet34': 'https://s3.amazonaws.com/pytorch/models/resnet34-333f7ec4.pth',
-  'resnet50': 'https://s3.amazonaws.com/pytorch/models/resnet50-19c8e357.pth',
-  'resnet101': 'https://s3.amazonaws.com/pytorch/models/resnet101-5d3b4d8f.pth',
-  'resnet152': 'https://s3.amazonaws.com/pytorch/models/resnet152-b121ed2d.pth',
-}
-
-
-def conv3x3(in_planes, out_planes, stride=1):
-  "3x3 convolution with padding"
-  return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
-           padding=1, bias=False)
-
-
-class BasicBlock(nn.Module):
-  expansion = 1
-
-  def __init__(self, inplanes, planes, stride=1, downsample=None):
-    super(BasicBlock, self).__init__()
-    self.conv1 = conv3x3(inplanes, planes, stride)
-    self.bn1 = nn.BatchNorm2d(planes)
-    self.relu = nn.ReLU(inplace=True)
-    self.conv2 = conv3x3(planes, planes)
-    self.bn2 = nn.BatchNorm2d(planes)
-    self.downsample = downsample
-    self.stride = stride
-
-  def forward(self, x):
-    residual = x
-
-    out = self.conv1(x)
-    out = self.bn1(out)
-    out = self.relu(out)
-
-    out = self.conv2(out)
-    out = self.bn2(out)
-
-    if self.downsample is not None:
-      residual = self.downsample(x)
-
-    out += residual
-    out = self.relu(out)
-
-    return out
-
-
-class Bottleneck(nn.Module):
-  expansion = 4
-
-  def __init__(self, inplanes, planes, stride=1, downsample=None):
-    super(Bottleneck, self).__init__()
-    self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, stride=stride, bias=False) # change
-    self.bn1 = nn.BatchNorm2d(planes)
-    self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, # change
-                 padding=1, bias=False)
-    self.bn2 = nn.BatchNorm2d(planes)
-    self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
-    self.bn3 = nn.BatchNorm2d(planes * 4)
-    self.relu = nn.ReLU(inplace=True)
-    self.downsample = downsample
-    self.stride = stride
-
-  def forward(self, x):
-    residual = x
-
-    out = self.conv1(x)
-    out = self.bn1(out)
-    out = self.relu(out)
-
-    out = self.conv2(out)
-    out = self.bn2(out)
-    out = self.relu(out)
-
-    out = self.conv3(out)
-    out = self.bn3(out)
-
-    if self.downsample is not None:
-      residual = self.downsample(x)
-
-    out += residual
-    out = self.relu(out)
-
-    return out
-
-
-class ResNet(nn.Module):
+class ResNet(torchvision.models.resnet.ResNet):
   def __init__(self, block, layers, num_classes=1000):
     self.inplanes = 64
-    super(ResNet, self).__init__()
-    self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
-                 bias=False)
-    self.bn1 = nn.BatchNorm2d(64)
-    self.relu = nn.ReLU(inplace=True)
-    # maxpool different from pytorch-resnet, to match tf-faster-rcnn
-    self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
-    self.layer1 = self._make_layer(block, 64, layers[0])
-    self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
-    self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+    super(ResNet, self).__init__(block, layers, num_classes)
+    # change to match the caffe resnet
+    for i in range(2, 4):
+      getattr(self, 'layer%d'%i)[0].conv1.stride = (2,2)
+      getattr(self, 'layer%d'%i)[0].conv2.stride = (1,1)
     # use stride 1 for the last conv4 layer (same as tf-faster-rcnn)
-    self.layer4 = self._make_layer(block, 512, layers[3], stride=1)
-
-    for m in self.modules():
-      if isinstance(m, nn.Conv2d):
-        n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
-        m.weight.data.normal_(0, math.sqrt(2. / n))
-      elif isinstance(m, nn.BatchNorm2d):
-        m.weight.data.fill_(1)
-        m.bias.data.zero_()
-
-  def _make_layer(self, block, planes, blocks, stride=1):
-    downsample = None
-    if stride != 1 or self.inplanes != planes * block.expansion:
-      downsample = nn.Sequential(
-        nn.Conv2d(self.inplanes, planes * block.expansion,
-              kernel_size=1, stride=stride, bias=False),
-        nn.BatchNorm2d(planes * block.expansion),
-      )
-
-    layers = []
-    layers.append(block(self.inplanes, planes, stride, downsample))
-    self.inplanes = planes * block.expansion
-    for i in range(1, blocks):
-      layers.append(block(self.inplanes, planes))
-
-    return nn.Sequential(*layers)
+    self.layer4[0].conv2.stride = (1,1)
+    self.layer4[0].downsample[0].stride = (1,1)
+
+    del self.avgpool, self.fc
+
 
 def resnet18(pretrained=False):
   """Constructs a ResNet-18 model.