Test ability to modify pooling module of network
(net, img_size)
| 81 | |
| 82 | @pytest.mark.parametrize('img_size', [224, 256, 512]) |
| 83 | def test_modify_pool(net, img_size): |
| 84 | """Test ability to modify pooling module of network""" |
| 85 | |
| 86 | class AdaptiveMaxAvgPool(nn.Module): |
| 87 | |
| 88 | def __init__(self): |
| 89 | super().__init__() |
| 90 | self.ada_avgpool = nn.AdaptiveAvgPool2d(1) |
| 91 | self.ada_maxpool = nn.AdaptiveMaxPool2d(1) |
| 92 | |
| 93 | def forward(self, x): |
| 94 | avg_x = self.ada_avgpool(x) |
| 95 | max_x = self.ada_maxpool(x) |
| 96 | x = torch.cat((avg_x, max_x), dim=1) |
| 97 | return x |
| 98 | |
| 99 | avg_pooling = AdaptiveMaxAvgPool() |
| 100 | fc = nn.Linear(net._fc.in_features * 2, net._global_params.num_classes) |
| 101 | |
| 102 | net._avg_pooling = avg_pooling |
| 103 | net._fc = fc |
| 104 | |
| 105 | data = torch.zeros((2, 3, img_size, img_size)) |
| 106 | output = net(data) |
| 107 | assert not torch.isnan(output).any() |
nothing calls this directly
no test coverage detected