| 212 | ) |
| 213 | |
| 214 | def make_encoder(self, in_channels, encoder_arch): |
| 215 | layers = [] |
| 216 | down_factor = 0 |
| 217 | for v in encoder_arch: |
| 218 | if v == 'MP': |
| 219 | layers += [nn.MaxPool2d(2, stride=2, ceil_mode=True)] |
| 220 | down_factor += 1 |
| 221 | else: |
| 222 | conv1 = nn.Conv2d(in_channels, v, kernel_size=5, padding=2) |
| 223 | conv2 = nn.Conv2d(v, v, kernel_size=5, padding=2) |
| 224 | |
| 225 | layers += [conv1, nn.InstanceNorm2d(v, affine=True), nn.ReLU(inplace=True), |
| 226 | conv2, nn.InstanceNorm2d(v, affine=True), nn.ReLU(inplace=True)] |
| 227 | in_channels = v |
| 228 | return nn.Sequential(*layers), 2 ** down_factor |
| 229 | |
| 230 | |
| 231 | def forward(self, image): |