diff options
author | Jordan Gong <jordan.gong@protonmail.com> | 2021-02-18 18:37:13 +0800 |
---|---|---|
committer | Jordan Gong <jordan.gong@protonmail.com> | 2021-02-18 18:37:13 +0800 |
commit | 5d4a7460597f2829fad04ab9a8157a0999904d44 (patch) | |
tree | ca02bc8c46241f39ce806c45d6e47a5672414c84 /models/layers.py | |
parent | 5097f32e39ca66723808c21288d0c8698e3e313c (diff) | |
parent | 84a3d5991f2f7272d1be54ad6cfe6ce695f915a0 (diff) |
Merge branch 'master' into data_parallel
Diffstat (limited to 'models/layers.py')
-rw-r--r-- | models/layers.py | 8 |
1 files 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 |