From 745d18a3249b79a3c91ff317c5109d08be6e4214 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Tue, 6 Jun 2023 22:18:19 +0200 Subject: [PATCH] add feature smoothing --- models/pose_encoder.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/models/pose_encoder.py b/models/pose_encoder.py index 232356a..da2702e 100644 --- a/models/pose_encoder.py +++ b/models/pose_encoder.py @@ -42,6 +42,11 @@ class ConvNeXtReverseMaskPose(nn.Module): if norm_layer is None: norm_layer = partial(LayerNorm2d, eps=1e-6) + stage: List[nn.Module] = [] + for _ in range(3): + stage.append(block(block_setting[0].input_channels, layer_scale, 0)) + self.stage_smooth = nn.Sequential(*stage).to(device) + self.first = nn.Sequential( nn.ConvTranspose2d(block_setting[0].input_channels, block_setting[0].out_channels, kernel_size=2, stride=2), norm_layer(block_setting[0].out_channels), @@ -111,7 +116,8 @@ class ConvNeXtReverseMaskPose(nn.Module): input, features = x else: input = x - x = self.first(input) + x = self.stage_smooth(input) + x = self.first(x) f = self.layers(x) x = self.stage_pose(f) pose = torch.sigmoid(self.last(x))