diff options
author | Jordan Gong <jordan.gong@protonmail.com> | 2021-02-21 19:01:20 +0800 |
---|---|---|
committer | Jordan Gong <jordan.gong@protonmail.com> | 2021-02-21 19:01:20 +0800 |
commit | 1014543e8b2cecbe6fdf1a0135bbefccee9c0d41 (patch) | |
tree | 90d726834bec29c8816800c71b75aad663edc858 /models/part_net.py | |
parent | c538919cb69e35a46811aef0b23baefe6a4c499c (diff) | |
parent | c52fdc2748e272a5195303299a9739291be32281 (diff) |
Merge branch 'master' into python3.8
# Conflicts:
# models/part_net.py
# models/rgb_part_net.py
Diffstat (limited to 'models/part_net.py')
-rw-r--r-- | models/part_net.py | 18 |
1 files changed, 4 insertions, 14 deletions
diff --git a/models/part_net.py b/models/part_net.py index f34f993..d9b954f 100644 --- a/models/part_net.py +++ b/models/part_net.py @@ -111,32 +111,22 @@ class TemporalFeatureAggregator(nn.Module): class PartNet(nn.Module): def __init__( self, - in_channels: int = 3, - feature_channels: int = 32, - kernel_sizes: Tuple[Tuple, ...] = ((5, 3), (3, 3), (3, 3)), - paddings: Tuple[Tuple, ...] = ((2, 1), (1, 1), (1, 1)), - halving: Tuple[int, ...] = (0, 2, 3), + in_channels: int = 128, squeeze_ratio: int = 4, num_part: int = 16 ): super().__init__() self.num_part = num_part - self.fpfe = FrameLevelPartFeatureExtractor( - in_channels, feature_channels, kernel_sizes, paddings, halving - ) - - num_fconv_blocks = len(self.fpfe.fconv_blocks) - self.tfa_in_channels = feature_channels * 2 ** (num_fconv_blocks - 1) self.tfa = TemporalFeatureAggregator( - self.tfa_in_channels, squeeze_ratio, self.num_part + in_channels, squeeze_ratio, self.num_part ) self.avg_pool = nn.AdaptiveAvgPool2d(1) self.max_pool = nn.AdaptiveMaxPool2d(1) def forward(self, x): - n, t, _, _, _ = x.size() - x = self.fpfe(x) + n, t, c, h, w = x.size() + x = x.view(n * t, c, h, w) # n * t x c x h x w # Horizontal Pooling |