| 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) |