MIMO-Unet 网络架构详解

def forward(self, x):#1 2 208 208
        x_2 = F.interpolate(x, scale_factor=0.5,  recompute_scale_factor=True) #1 3 104 104
        x_4 = F.interpolate(x_2, scale_factor=0.5,  recompute_scale_factor=True) # 1 3 52 52
        z2 = self.SCM2(x_2) #1 64 104 104
        z4 = self.SCM1(x_4) #1 128 52 52

        outputs = list()

        x_ = self.feat_extract[0](x) # 1 32 208 208
        res1 = self.Encoder[0](x_) # 1 32 208 208

        z = self.feat_extract[1](res1) # 1 64 104 104
        z = self.FAM2(z, z2) # 1 64 104 104
        res2 = self.Encoder[1](z) # 1 64 104 104

        z = self.feat_extract[2](res2) #1 128 52 52
        z = self.FAM1(z, z4) #1 128 52 52
        z = self.Encoder[2](z) #1 128 52 52

        z12 = F.interpolate(res1, scale_factor=0.5,  recompute_scale_factor=True) #1 32 104 104
        z21 = F.interpolate(res2, scale_factor=2,  recompute_scale_factor=True) #1 64 208 208
        z42 = F.interpolate(z, scale_factor=2,  recompute_scale_factor=True) # 1 128 104 104
        z41 = F.interpolate(z42, scale_factor=2,  recompute_scale_factor=True) #1 128 208 208

        res2 = self.AFFs[1](z12, res2, z42) # 1 64 104 104
        res1 = self.AFFs[0](res1, z21, z41) # 1 32 208 208

        z = self.Decoder[0](z) #1 128 52 52
        z_ = self.ConvsOut[0](z) #1 3 52 52
        z = self.feat_extract[3](z) #1 64 104 104
        outputs.append(z_+x_4)

        z = torch.cat([z, res2], dim=1)# 1 128 104 104
        z = self.Convs[0](z) # 1 64 104 104
        z = self.Decoder[1](z) # 1 64 104 104
        z_ = self.ConvsOut[1](z) # 1 3 104 104
        z = self.feat_extract[4](z)
        outputs.append(z_+x_2)

        z = torch.cat([z, res1], dim=1)# 1 64 208 208
        z = self.Convs[1](z) # 1 32 208 208
        z = self.Decoder[2](z) # 1 32 208 208
        z = self.feat_extract[5](z)# 1 3 208 208
        outputs.append(z+x)

        return outputs