Added test for modifying pooling in final layers of network

Author: khornlund

efficientnet_pytorch/model.py

@@ -174,13 +174,13 @@ def extract_features(self, inputs):
 
     def forward(self, inputs):
         """ Calls extract_features to extract features, applies final linear layer, and returns logits. """
-
+        bs = inputs.size(0)
         # Convolution layers
         x = self.extract_features(inputs)
 
         # Pooling and final linear layer
         x = self._avg_pooling(x)
-        x = x.view(x.size(0), -1)
+        x = x.view(bs, -1)
         x = self._dropout(x)
         x = self._fc(x)
         return x

tests/test_model.py

@@ -19,7 +19,7 @@ def pretrained(request):
     return request.param
 
 
-@pytest.fixture(scope='module')
+@pytest.fixture(scope='function')
 def net(model, pretrained):
     return EfficientNet.from_pretrained(model) if pretrained else EfficientNet.from_name(model)
 
@@ -59,8 +59,8 @@ def test_dropout_update(net):
 
 
 @pytest.mark.parametrize('img_size', [224, 256, 512])
-def test_modify_head(net, img_size):
-    """Test ability to modify final layers of network"""
+def test_modify_dropout(net, img_size):
+    """Test ability to modify dropout and fc modules of network"""
     dropout = nn.Sequential(OrderedDict([
         ('_bn2', nn.BatchNorm1d(net._bn1.num_features)),
         ('_drop1', nn.Dropout(p=net._global_params.dropout_rate)),
@@ -77,3 +77,32 @@ def test_modify_head(net, img_size):
     data = torch.zeros((2, 3, img_size, img_size))
     output = net(data)
     assert not torch.isnan(output).any()
+
+
+@pytest.mark.parametrize('img_size', [224, 256, 512])
+def test_modify_norm(net, img_size):
+    """Test ability to modify norm layer of network"""
+
+    class AdaptiveMaxAvgPool(nn.Module):
+
+        def __init__(self):
+            super().__init__()
+            self.ada_avgpool = nn.AdaptiveAvgPool2d(1)
+            self.ada_maxpool = nn.AdaptiveMaxPool2d(1)
+
+        def forward(self, x):
+            avg_x = self.ada_avgpool(x)
+            max_x = self.ada_maxpool(x)
+            x = torch.cat((avg_x, max_x), dim=1)
+            x = x.view(x.size(0), -1)
+            return x
+
+    avg_pooling = AdaptiveMaxAvgPool()
+    fc = nn.Linear(net._fc.in_features * 2, net._global_params.num_classes)
+
+    net._avg_pooling = avg_pooling
+    net._fc = fc
+
+    data = torch.zeros((2, 3, img_size, img_size))
+    output = net(data)
+    assert not torch.isnan(output).any()