diff options
Diffstat (limited to 'models/hpm.py')
| -rw-r--r-- | models/hpm.py | 5 |
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 |