| 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) |
no outgoing calls