| 236 | |
| 237 | |
| 238 | class AO(nn.Module): |
| 239 | # Adaptive output module |
| 240 | def __init__(self, inchannels, outchannels, upfactor=2): |
| 241 | super(AO, self).__init__() |
| 242 | self.inchannels = inchannels |
| 243 | self.outchannels = outchannels |
| 244 | self.upfactor = upfactor |
| 245 | |
| 246 | self.adapt_conv = nn.Sequential( |
| 247 | nn.Conv2d(in_channels=self.inchannels, out_channels=self.inchannels // 2, kernel_size=3, padding=1, |
| 248 | stride=1, bias=True), \ |
| 249 | nn.BatchNorm2d(num_features=self.inchannels // 2), \ |
| 250 | nn.ReLU(inplace=True), \ |
| 251 | nn.Conv2d(in_channels=self.inchannels // 2, out_channels=self.outchannels, kernel_size=3, padding=1, |
| 252 | stride=1, bias=True), \ |
| 253 | nn.Upsample(scale_factor=self.upfactor, mode='bilinear', align_corners=True)) |
| 254 | |
| 255 | self.init_params() |
| 256 | |
| 257 | def forward(self, x): |
| 258 | x = self.adapt_conv(x) |
| 259 | return x |
| 260 | |
| 261 | def init_params(self): |
| 262 | for m in self.modules(): |
| 263 | if isinstance(m, nn.Conv2d): |
| 264 | # init.kaiming_normal_(m.weight, mode='fan_out') |
| 265 | init.normal_(m.weight, std=0.01) |
| 266 | # init.xavier_normal_(m.weight) |
| 267 | if m.bias is not None: |
| 268 | init.constant_(m.bias, 0) |
| 269 | elif isinstance(m, nn.ConvTranspose2d): |
| 270 | # init.kaiming_normal_(m.weight, mode='fan_out') |
| 271 | init.normal_(m.weight, std=0.01) |
| 272 | # init.xavier_normal_(m.weight) |
| 273 | if m.bias is not None: |
| 274 | init.constant_(m.bias, 0) |
| 275 | elif isinstance(m, nn.BatchNorm2d): # NN.Batchnorm2d |
| 276 | init.constant_(m.weight, 1) |
| 277 | init.constant_(m.bias, 0) |
| 278 | elif isinstance(m, nn.Linear): |
| 279 | init.normal_(m.weight, std=0.01) |
| 280 | if m.bias is not None: |
| 281 | init.constant_(m.bias, 0) |
| 282 | |
| 283 | |
| 284 | |