adjust network to match TensorFlow

Author: TheCodez

docs/source/models.rst

@@ -150,5 +150,4 @@ GoogLeNet
 ------------
 
 .. autofunction:: googlenet
-.. autofunction:: googlenet_bn
 

torchvision/models/googlenet.py

@@ -3,11 +3,11 @@
 import torch.nn.functional as F
 from torch.utils import model_zoo
 
-__all__ = ['GoogLeNet', 'googlenet', 'googlenet_bn']
+__all__ = ['GoogLeNet', 'googlenet']
 
 model_urls = {
-    'googlenet': '',
-    'googlenet_bn': ''
+    # GoogLeNet ported from TensorFlow
+    'googlenet': 'https://github.com/TheCodez/vision/releases/download/1.0/googlenet-1378be20.pth',
 }
 
 
@@ -18,6 +18,8 @@ def googlenet(pretrained=False, **kwargs):
         pretrained (bool): If True, returns a model pre-trained on ImageNet
     """
     if pretrained:
+        if 'transform_input' not in kwargs:
+            kwargs['transform_input'] = True
         kwargs['init_weights'] = False
         model = GoogLeNet(**kwargs)
         model.load_state_dict(model_zoo.load_url(model_urls['googlenet']))
@@ -26,51 +28,35 @@ def googlenet(pretrained=False, **kwargs):
     return GoogLeNet(**kwargs)
 
 
-def googlenet_bn(pretrained=False, **kwargs):
-    r"""GoogLeNet (Inception v1) model architecture with batch normalization from
-    `"Going Deeper with Convolutions" <http://arxiv.org/abs/1409.4842>`_.
-    Args:
-        pretrained (bool): If True, returns a model pre-trained on ImageNet
-    """
-    if pretrained:
-        kwargs['init_weights'] = False
-        model = GoogLeNet(batch_norm=True, **kwargs)
-        model.load_state_dict(model_zoo.load_url(model_urls['googlenet_bn']))
-        return model
-
-    return GoogLeNet(batch_norm=True, **kwargs)
-
-
 class GoogLeNet(nn.Module):
 
-    def __init__(self, num_classes=1000, aux_logits=True, batch_norm=False, init_weights=True):
+    def __init__(self, num_classes=1000, aux_logits=True, transform_input=False, init_weights=True):
         super(GoogLeNet, self).__init__()
         self.aux_logits = aux_logits
+        self.transform_input = transform_input
 
-        self.conv1 = BasicConv2d(3, 64, batch_norm, kernel_size=7, stride=2, padding=3)
+        self.conv1 = BasicConv2d(3, 64, kernel_size=7, stride=2, padding=3)
         self.maxpool1 = nn.MaxPool2d(3, stride=2, ceil_mode=True)
-        self.lrn1 = nn.LocalResponseNorm(5, alpha=0.0001, beta=0.75)
-        self.conv2 = BasicConv2d(64, 64, batch_norm, kernel_size=1)
-        self.conv3 = BasicConv2d(64, 192, batch_norm, kernel_size=3, padding=1)
-        self.lrn2 = nn.LocalResponseNorm(5, alpha=0.0001, beta=0.75)
+        self.conv2 = BasicConv2d(64, 64, kernel_size=1)
+        self.conv3 = BasicConv2d(64, 192, kernel_size=3, padding=1)
         self.maxpool2 = nn.MaxPool2d(3, stride=2, ceil_mode=True)
 
-        self.inception3a = Inception(192, 64, 96, 128, 16, 32, 32, batch_norm)
-        self.inception3b = Inception(256, 128, 128, 192, 32, 96, 64, batch_norm)
+        self.inception3a = Inception(192, 64, 96, 128, 16, 32, 32)
+        self.inception3b = Inception(256, 128, 128, 192, 32, 96, 64)
         self.maxpool3 = nn.MaxPool2d(3, stride=2, ceil_mode=True)
 
-        self.inception4a = Inception(480, 192, 96, 208, 16, 48, 64, batch_norm)
-        self.inception4b = Inception(512, 160, 112, 224, 24, 64, 64, batch_norm)
-        self.inception4c = Inception(512, 128, 128, 256, 24, 64, 64, batch_norm)
-        self.inception4d = Inception(512, 112, 144, 288, 32, 64, 64, batch_norm)
-        self.inception4e = Inception(528, 256, 160, 320, 32, 128, 128, batch_norm)
+        self.inception4a = Inception(480, 192, 96, 208, 16, 48, 64)
+        self.inception4b = Inception(512, 160, 112, 224, 24, 64, 64)
+        self.inception4c = Inception(512, 128, 128, 256, 24, 64, 64)
+        self.inception4d = Inception(512, 112, 144, 288, 32, 64, 64)
+        self.inception4e = Inception(528, 256, 160, 320, 32, 128, 128)
         self.maxpool4 = nn.MaxPool2d(3, stride=2, ceil_mode=True)
 
-        self.inception5a = Inception(832, 256, 160, 320, 32, 128, 128, batch_norm)
-        self.inception5b = Inception(832, 384, 192, 384, 48, 128, 128, batch_norm)
+        self.inception5a = Inception(832, 256, 160, 320, 32, 128, 128)
+        self.inception5b = Inception(832, 384, 192, 384, 48, 128, 128)
         if aux_logits:
-            self.aux1 = InceptionAux(512, num_classes, batch_norm)
-            self.aux2 = InceptionAux(528, num_classes, batch_norm)
+            self.aux1 = InceptionAux(512, num_classes)
+            self.aux2 = InceptionAux(528, num_classes)
         self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
         self.dropout = nn.Dropout(0.4)
         self.fc = nn.Linear(1024, num_classes)
@@ -92,12 +78,16 @@ def _initialize_weights(self):
                 nn.init.constant_(m.bias, 0)
 
     def forward(self, x):
+        if self.transform_input:
+            x_ch0 = torch.unsqueeze(x[:, 0], 1) * (0.229 / 0.5) + (0.485 - 0.5) / 0.5
+            x_ch1 = torch.unsqueeze(x[:, 1], 1) * (0.224 / 0.5) + (0.456 - 0.5) / 0.5
+            x_ch2 = torch.unsqueeze(x[:, 2], 1) * (0.225 / 0.5) + (0.406 - 0.5) / 0.5
+            x = torch.cat((x_ch0, x_ch1, x_ch2), 1)
+
         x = self.conv1(x)
         x = self.maxpool1(x)
-        x = self.lrn1(x)
         x = self.conv2(x)
         x = self.conv3(x)
-        x = self.lrn2(x)
         x = self.maxpool2(x)
 
         x = self.inception3a(x)
@@ -129,24 +119,24 @@ def forward(self, x):
 
 class Inception(nn.Module):
 
-    def __init__(self, in_channels, ch1x1, ch3x3red, ch3x3, ch5x5red, ch5x5, pool_proj, batch_norm=False):
+    def __init__(self, in_channels, ch1x1, ch3x3red, ch3x3, ch5x5red, ch5x5, pool_proj):
         super(Inception, self).__init__()
 
-        self.branch1 = BasicConv2d(in_channels, ch1x1, batch_norm, kernel_size=1)
+        self.branch1 = BasicConv2d(in_channels, ch1x1, kernel_size=1)
 
         self.branch2 = nn.Sequential(
-            BasicConv2d(in_channels, ch3x3red, batch_norm, kernel_size=1),
-            BasicConv2d(ch3x3red, ch3x3, batch_norm, kernel_size=3, padding=1)
+            BasicConv2d(in_channels, ch3x3red, kernel_size=1),
+            BasicConv2d(ch3x3red, ch3x3, kernel_size=3, padding=1)
         )
 
         self.branch3 = nn.Sequential(
-            BasicConv2d(in_channels, ch5x5red, batch_norm, kernel_size=1),
-            BasicConv2d(ch5x5red, ch5x5, batch_norm, kernel_size=5, padding=2)
+            BasicConv2d(in_channels, ch5x5red, kernel_size=1),
+            BasicConv2d(ch5x5red, ch5x5, kernel_size=3, padding=1)
         )
 
         self.branch4 = nn.Sequential(
             nn.MaxPool2d(kernel_size=3, stride=1, padding=1, ceil_mode=True),
-            BasicConv2d(in_channels, pool_proj, batch_norm, kernel_size=1)
+            BasicConv2d(in_channels, pool_proj, kernel_size=1)
         )
 
     def forward(self, x):
@@ -161,11 +151,11 @@ def forward(self, x):
 
 class InceptionAux(nn.Module):
 
-    def __init__(self, in_channels, num_classes, batch_norm=False):
+    def __init__(self, in_channels, num_classes):
         super(InceptionAux, self).__init__()
-        self.conv = BasicConv2d(in_channels, 128, batch_norm, kernel_size=1)
+        self.conv = BasicConv2d(in_channels, 128, kernel_size=1)
 
-        self.fc1 = nn.Linear(128 * 4 * 4, 1024)
+        self.fc1 = nn.Linear(2048, 1024)
         self.fc2 = nn.Linear(1024, num_classes)
 
     def forward(self, x):
@@ -182,18 +172,12 @@ def forward(self, x):
 
 class BasicConv2d(nn.Module):
 
-    def __init__(self, in_channels, out_channels, batch_norm=False, **kwargs):
+    def __init__(self, in_channels, out_channels, **kwargs):
         super(BasicConv2d, self).__init__()
-        self.batch_norm = batch_norm
-
-        if self.batch_norm:
-            self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs)
-            self.bn = nn.BatchNorm2d(out_channels, eps=0.001)
-        else:
-            self.conv = nn.Conv2d(in_channels, out_channels, **kwargs)
+        self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs)
+        self.bn = nn.BatchNorm2d(out_channels, eps=0.001)
 
     def forward(self, x):
         x = self.conv(x)
-        if self.batch_norm:
-            x = self.bn(x)
+        x = self.bn(x)
         return F.relu(x, inplace=True)