From e6703e1fff20975bada71ce12644e0c8e58c1922 Mon Sep 17 00:00:00 2001 From: Jordan Gong Date: Thu, 18 Feb 2021 17:15:56 +0800 Subject: Remove 1x1 conv layers when not used --- models/layers.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/models/layers.py b/models/layers.py index 7b6ba5c..ef53a95 100644 --- a/models/layers.py +++ b/models/layers.py @@ -162,7 +162,7 @@ class BasicConv1d(nn.Module): return self.conv(x) -class HorizontalPyramidPooling(BasicConv2d): +class HorizontalPyramidPooling(nn.Module): def __init__( self, in_channels: int, @@ -172,8 +172,10 @@ class HorizontalPyramidPooling(BasicConv2d): use_max_pool: bool = False, **kwargs ): - super().__init__(in_channels, out_channels, kernel_size=1, **kwargs) + super().__init__() self.use_1x1conv = use_1x1conv + if use_1x1conv: + self.conv = BasicConv2d(in_channels, out_channels, 1, **kwargs) 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.' @@ -188,5 +190,5 @@ class HorizontalPyramidPooling(BasicConv2d): elif not self.use_avg_pool and self.use_max_pool: x = self.max_pool(x) if self.use_1x1conv: - x = super().forward(x) + x = self.conv(x) return x -- cgit v1.2.3