From 63a12171b90ef6869cf0faae5cca55c738c3108f Mon Sep 17 00:00:00 2001 From: Karl Hornlund Date: Tue, 3 Sep 2019 12:53:10 +0930 Subject: [PATCH 1/5] Updated model to use and simplified implementation --- efficientnet_pytorch/model.py | 9 ++-- tests/test_model.py | 79 +++++++++++++++++++++++++++++++++++ 2 files changed, 84 insertions(+), 4 deletions(-) create mode 100644 tests/test_model.py diff --git a/efficientnet_pytorch/model.py b/efficientnet_pytorch/model.py index 3a42217..c139c0c 100644 --- a/efficientnet_pytorch/model.py +++ b/efficientnet_pytorch/model.py @@ -150,7 +150,8 @@ def __init__(self, blocks_args=None, global_params=None): self._bn1 = nn.BatchNorm2d(num_features=out_channels, momentum=bn_mom, eps=bn_eps) # Final linear layer - self._dropout = self._global_params.dropout_rate + self._avg_pooling = nn.AdaptiveAvgPool2d(1) + self._dropout = nn.Dropout(self._global_params.dropout_rate) self._fc = nn.Linear(out_channels, self._global_params.num_classes) def extract_features(self, inputs): @@ -178,9 +179,9 @@ def forward(self, inputs): x = self.extract_features(inputs) # Pooling and final linear layer - x = F.adaptive_avg_pool2d(x, 1).squeeze(-1).squeeze(-1) - if self._dropout: - x = F.dropout(x, p=self._dropout, training=self.training) + x = self._avg_pooling(x) + x = x.view(x.size(0), -1) + x = self._dropout(x) x = self._fc(x) return x diff --git a/tests/test_model.py b/tests/test_model.py new file mode 100644 index 0000000..a778931 --- /dev/null +++ b/tests/test_model.py @@ -0,0 +1,79 @@ +from collections import OrderedDict + +import pytest +import torch +import torch.nn as nn + +from efficientnet_pytorch import EfficientNet + + +# -- fixtures ------------------------------------------------------------------------------------- + +@pytest.fixture(scope='module', params=[x for x in range(8)]) +def model(request): + return 'efficientnet-b{}'.format(request.param) + + +@pytest.fixture(scope='module', params=[True, False]) +def pretrained(request): + return request.param + + +@pytest.fixture(scope='module') +def net(model, pretrained): + return EfficientNet.from_pretrained(model) if pretrained else EfficientNet.from_name(model) + + +# -- tests ---------------------------------------------------------------------------------------- + +@pytest.mark.parametrize('img_size', [224, 256, 512]) +def test_forward(net, img_size): + """Test `.forward()` doesn't throw an error""" + data = torch.zeros((1, 3, img_size, img_size)) + output = net(data) + assert not torch.isnan(output).any() + + +def test_dropout_training(net): + """Test dropout `.training` is set by `.train()` on parent `nn.module`""" + net.train() + assert net._dropout.training == True + + +def test_dropout_eval(net): + """Test dropout `.training` is set by `.eval()` on parent `nn.module`""" + net.eval() + assert net._dropout.training == False + + +def test_dropout_update(net): + """Test dropout `.training` is updated by `.train()` and `.eval()` on parent `nn.module`""" + net.train() + assert net._dropout.training == True + net.eval() + assert net._dropout.training == False + net.train() + assert net._dropout.training == True + net.eval() + assert net._dropout.training == False + + +@pytest.mark.parametrize('img_size', [224, 256, 512]) +def test_modify_head(net, img_size): + """Test ability to modify final layers of network""" + dropout = nn.Sequential(OrderedDict([ + ('_bn2', nn.BatchNorm1d(net._bn1.num_features)), + ('_drop1', nn.Dropout(p=net._global_params.dropout_rate)), + ('_linear1', nn.Linear(net._bn1.num_features, 512)), + ('_relu', nn.ReLU()), + ('_bn3', nn.BatchNorm1d(512)), + ('_drop2', nn.Dropout(p=net._global_params.dropout_rate / 2)) + ])) + fc = nn.Linear(512, net._global_params.num_classes) + + net._dropout = dropout + net._fc = fc + + data = torch.zeros((2, 3, img_size, img_size)) + output = net(data) + assert not torch.isnan(output).any() From 4f1cf2fa7460c2f921931f9f292fede36e5b6584 Mon Sep 17 00:00:00 2001 From: Karl Hornlund Date: Tue, 3 Sep 2019 12:56:17 +0930 Subject: [PATCH 2/5] Updated fixture to only test models 0-3 --- tests/test_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_model.py b/tests/test_model.py index a778931..32695f0 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -9,7 +9,7 @@ # -- fixtures ------------------------------------------------------------------------------------- -@pytest.fixture(scope='module', params=[x for x in range(8)]) +@pytest.fixture(scope='module', params=[x for x in range(4)]) def model(request): return 'efficientnet-b{}'.format(request.param) From f938df23aa979d039f720e8c215f94bc1bc2ae53 Mon Sep 17 00:00:00 2001 From: Karl Hornlund Date: Tue, 3 Sep 2019 14:45:07 +0930 Subject: [PATCH 3/5] Added test for modifying pooling in final layers of network --- efficientnet_pytorch/model.py | 4 ++-- tests/test_model.py | 35 ++++++++++++++++++++++++++++++++--- 2 files changed, 34 insertions(+), 5 deletions(-) diff --git a/efficientnet_pytorch/model.py b/efficientnet_pytorch/model.py index c139c0c..5021b15 100644 --- a/efficientnet_pytorch/model.py +++ b/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 diff --git a/tests/test_model.py b/tests/test_model.py index 32695f0..d99ec34 100644 --- a/tests/test_model.py +++ b/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() \ No newline at end of file From 7fa4d0456466e0a9569bc5c14f8ecd67ee5fbee8 Mon Sep 17 00:00:00 2001 From: Karl Hornlund Date: Tue, 3 Sep 2019 16:13:52 +0930 Subject: [PATCH 4/5] Removed reshape in --- tests/test_model.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_model.py b/tests/test_model.py index d99ec34..4b060b1 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -94,7 +94,6 @@ 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() From d4aed289a36441778ccb869e7bd2b8eda35cc8b2 Mon Sep 17 00:00:00 2001 From: Karl Hornlund Date: Tue, 3 Sep 2019 17:11:48 +0930 Subject: [PATCH 5/5] Updated comments --- tests/test_model.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_model.py b/tests/test_model.py index 4b060b1..a138a2b 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -80,8 +80,8 @@ def test_modify_dropout(net, img_size): @pytest.mark.parametrize('img_size', [224, 256, 512]) -def test_modify_norm(net, img_size): - """Test ability to modify norm layer of network""" +def test_modify_pool(net, img_size): + """Test ability to modify pooling module of network""" class AdaptiveMaxAvgPool(nn.Module): @@ -104,4 +104,4 @@ def forward(self, x): data = torch.zeros((2, 3, img_size, img_size)) output = net(data) - assert not torch.isnan(output).any() \ No newline at end of file + assert not torch.isnan(output).any()