Hourglass Encoder
| 184 | |
| 185 | |
| 186 | class Encoder(nn.Module): |
| 187 | """ |
| 188 | Hourglass Encoder |
| 189 | """ |
| 190 | |
| 191 | def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256): |
| 192 | super(Encoder, self).__init__() |
| 193 | |
| 194 | down_blocks = [] |
| 195 | for i in range(num_blocks): |
| 196 | down_blocks.append(DownBlock3d(in_features if i == 0 else min(max_features, block_expansion * (2 ** i)), min(max_features, block_expansion * (2 ** (i + 1))), kernel_size=3, padding=1)) |
| 197 | self.down_blocks = nn.ModuleList(down_blocks) |
| 198 | |
| 199 | def forward(self, x): |
| 200 | outs = [x] |
| 201 | for down_block in self.down_blocks: |
| 202 | outs.append(down_block(outs[-1])) |
| 203 | return outs |
| 204 | |
| 205 | |
| 206 | class Decoder(nn.Module): |