Add pooling options in HPM

According to [1], we can use GAP and GMP together, or one of both in ablation study.

[1]Y. Fu et al., “Horizontal pyramid matching for person re-identification,” in Proceedings of the AAAI Conference on Artificial Intelligence, 2019, vol. 33, pp. 8295–8302.

1 files changed, 13 insertions, 6 deletions

diff --git a/models/layers.py b/models/layers.py
index f824078..62a3cc6 100644
--- a/models/layers.py
+++ b/models/layers.py
@@ -141,16 +141,23 @@ class HorizontalPyramidPooling(BasicConv2d):
             in_channels: int,
             out_channels: int,
             kernel_size: Union[int, tuple[int, int]] = 1,
-            use_avg_pool: bool = False,
+            use_avg_pool: bool = True,
+            use_max_pool: bool = True,
             **kwargs
     ):
         super().__init__(in_channels, out_channels, kernel_size, **kwargs)
-        if use_avg_pool:
-            self.pool = nn.AdaptiveAvgPool2d(1)
-        else:
-            self.pool = nn.AdaptiveMaxPool2d(1)
+        self.use_avg_pool = use_avg_pool
+        self.use_max_pool = use_max_pool
+        assert use_avg_pool or use_max_pool, 'Pooling layer(s) required.'
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.max_pool = nn.AdaptiveMaxPool2d(1)
 
     def forward(self, x):
-        x = self.pool(x)
+        if self.use_avg_pool and self.use_max_pool:
+            x = self.avg_pool(x) + self.max_pool(x)
+        elif self.use_avg_pool and not self.use_max_pool:
+            x = self.avg_pool(x)
+        elif not self.use_avg_pool and self.use_max_pool:
+            x = self.max_pool(x)
         x = super().forward(x)
         return x