Hourglass Decoder
| 204 | |
| 205 | |
| 206 | class Decoder(nn.Module): |
| 207 | """ |
| 208 | Hourglass Decoder |
| 209 | """ |
| 210 | |
| 211 | def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256): |
| 212 | super(Decoder, self).__init__() |
| 213 | |
| 214 | up_blocks = [] |
| 215 | |
| 216 | for i in range(num_blocks)[::-1]: |
| 217 | in_filters = (1 if i == num_blocks - 1 else 2) * min(max_features, block_expansion * (2 ** (i + 1))) |
| 218 | out_filters = min(max_features, block_expansion * (2 ** i)) |
| 219 | up_blocks.append(UpBlock3d(in_filters, out_filters, kernel_size=3, padding=1)) |
| 220 | |
| 221 | self.up_blocks = nn.ModuleList(up_blocks) |
| 222 | self.out_filters = block_expansion + in_features |
| 223 | |
| 224 | self.conv = nn.Conv3d(in_channels=self.out_filters, out_channels=self.out_filters, kernel_size=3, padding=1) |
| 225 | self.norm = nn.BatchNorm3d(self.out_filters, affine=True) |
| 226 | |
| 227 | def forward(self, x): |
| 228 | out = x.pop() |
| 229 | for up_block in self.up_blocks: |
| 230 | out = up_block(out) |
| 231 | skip = x.pop() |
| 232 | out = torch.cat([out, skip], dim=1) |
| 233 | out = self.conv(out) |
| 234 | out = self.norm(out) |
| 235 | out = F.relu(out) |
| 236 | return out |
| 237 | |
| 238 | |
| 239 | class Hourglass(nn.Module): |