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, 4 insertions, 1 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