add feature smoothing

This commit is contained in:
2023-06-06 22:18:19 +02:00
parent cac382ef19
commit 745d18a324
+7 -1
View File
@@ -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))