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.

2 files changed, 17 insertions, 7 deletions

diff --git a/models/hpm.py b/models/hpm.py
index 1773f56..85a4e58 100644
--- a/models/hpm.py
+++ b/models/hpm.py
@@ -10,13 +10,15 @@ class HorizontalPyramidMatching(nn.Module):
             self,
             scales: tuple[int, ...] = (1, 2, 4, 8),
             out_channels: int = 256,
-            use_avg_pool: bool = False,
+            use_avg_pool: bool = True,
+            use_max_pool: bool = True,
             **kwargs
     ):
         super().__init__()
         self.scales = scales
         self.out_channels = out_channels
         self.use_avg_pool = use_avg_pool
+        self.use_max_pool = use_max_pool
 
         self.backbone = resnet50(pretrained=True)
         self.in_channels = self.backbone.layer4[-1].conv1.in_channels
@@ -29,6 +31,7 @@ class HorizontalPyramidMatching(nn.Module):
         pyramid = [HorizontalPyramidPooling(self.in_channels,
                                             self.out_channels,
                                             use_avg_pool=self.use_avg_pool,
+                                            use_max_pool=self.use_max_pool,
                                             **kwargs)
                    for _ in range(scale)]
         return pyramid
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