| 189 | |
| 190 | |
| 191 | class FFM(nn.Module): |
| 192 | def __init__(self, inchannels, midchannels, outchannels, upfactor=2): |
| 193 | super(FFM, self).__init__() |
| 194 | self.inchannels = inchannels |
| 195 | self.midchannels = midchannels |
| 196 | self.outchannels = outchannels |
| 197 | self.upfactor = upfactor |
| 198 | |
| 199 | self.ftb1 = FTB(inchannels=self.inchannels, midchannels=self.midchannels) |
| 200 | # self.ata = ATA(inchannels = self.midchannels) |
| 201 | self.ftb2 = FTB(inchannels=self.midchannels, midchannels=self.outchannels) |
| 202 | |
| 203 | self.upsample = nn.Upsample(scale_factor=self.upfactor, mode='bilinear', align_corners=True) |
| 204 | |
| 205 | self.init_params() |
| 206 | |
| 207 | def forward(self, low_x, high_x): |
| 208 | x = self.ftb1(low_x) |
| 209 | x = x + high_x |
| 210 | x = self.ftb2(x) |
| 211 | x = self.upsample(x) |
| 212 | |
| 213 | return x |
| 214 | |
| 215 | def init_params(self): |
| 216 | for m in self.modules(): |
| 217 | if isinstance(m, nn.Conv2d): |
| 218 | # init.kaiming_normal_(m.weight, mode='fan_out') |
| 219 | init.normal_(m.weight, std=0.01) |
| 220 | # init.xavier_normal_(m.weight) |
| 221 | if m.bias is not None: |
| 222 | init.constant_(m.bias, 0) |
| 223 | elif isinstance(m, nn.ConvTranspose2d): |
| 224 | # init.kaiming_normal_(m.weight, mode='fan_out') |
| 225 | init.normal_(m.weight, std=0.01) |
| 226 | # init.xavier_normal_(m.weight) |
| 227 | if m.bias is not None: |
| 228 | init.constant_(m.bias, 0) |
| 229 | elif isinstance(m, nn.BatchNorm2d): # NN.Batchnorm2d |
| 230 | init.constant_(m.weight, 1) |
| 231 | init.constant_(m.bias, 0) |
| 232 | elif isinstance(m, nn.Linear): |
| 233 | init.normal_(m.weight, std=0.01) |
| 234 | if m.bias is not None: |
| 235 | init.constant_(m.bias, 0) |
| 236 | |
| 237 | |
| 238 | class AO(nn.Module): |