| 58 | |
| 59 | |
| 60 | class FeatureExtractor(nn.Module): |
| 61 | def __init__(self, hint_channels, dim=2): |
| 62 | super(FeatureExtractor, self).__init__() |
| 63 | self.initial_conv = conv_nd(dim, hint_channels, 16, kernel_size=3, stride=1, padding=1) |
| 64 | self.layer1 = ResidualBlock(dim, 16, 16) |
| 65 | self.conv_before_res2 = conv_nd(dim, 16, 32, kernel_size=3, stride=1, padding=1) |
| 66 | self.layer2 = ResidualBlock(dim, 32, 32) |
| 67 | self.patch_merge1 = PatchMerging(patch_dim=32) |
| 68 | self.conv_before_res3 = conv_nd(dim, 64, 64, kernel_size=3, stride=1, padding=1) |
| 69 | self.layer3 = ResidualBlock(dim, 64, 64) |
| 70 | self.patch_merge2 = PatchMerging(patch_dim=64) |
| 71 | self.conv_before_res4 = conv_nd(dim, 128, 128, kernel_size=3, stride=1, padding=1) |
| 72 | self.layer4 = ResidualBlock(dim, 128, 128) |
| 73 | self.patch_merge3 = PatchMerging(patch_dim=128) |
| 74 | self.conv_before_res5 = conv_nd(dim, 256, 256, kernel_size=3, stride=1, padding=1) |
| 75 | self.layer5 = ResidualBlock(dim, 256, 256) |
| 76 | |
| 77 | def forward(self, x): |
| 78 | x = self.initial_conv(x) |
| 79 | x = self.layer1(x) |
| 80 | x = self.conv_before_res2(x) |
| 81 | x = self.layer2(x) |
| 82 | x = x.permute(0, 2, 3, 1) |
| 83 | x = self.patch_merge1(x) |
| 84 | x = x.permute(0, 3, 1, 2) |
| 85 | x = self.conv_before_res3(x) |
| 86 | x = self.layer3(x) |
| 87 | x = x.permute(0, 2, 3, 1) |
| 88 | x = self.patch_merge2(x) |
| 89 | x = x.permute(0, 3, 1, 2) |
| 90 | x = self.conv_before_res4(x) |
| 91 | x = self.layer4(x) |
| 92 | x = x.permute(0, 2, 3, 1) |
| 93 | x = self.patch_merge3(x) |
| 94 | x = x.permute(0, 3, 1, 2) |
| 95 | x = self.conv_before_res5(x) |
| 96 | x = self.layer5(x) |
| 97 | return x |