From 0178b7722c63820d701b0442b5946d9a042074d7 Mon Sep 17 00:00:00 2001 From: Jordan Gong Date: Wed, 30 Dec 2020 11:14:26 +0800 Subject: Add pooling options in HPM MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- models/hpm.py | 5 ++++- models/layers.py | 19 +++++++++++++------ 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 -- cgit v1.2.3