collaborative diagnosis-synthesis framework（CDSF）

论文阅读及简单的模型代码复现—— collaborative diagnosis-synthesis framework（CDSF）

题目：Collaborative Image Synthesis and Disease Diagnosis for Classification of Neurodegenerative Disorders with Incomplete Multi-modal Neuroimages

地址：Collaborative Image Synthesis and Disease Diagnosis for Classification of Neurodegenerative Disorders with Incomplete Multi-modal Neuroimages | SpringerLink

来源：MICCAI2021

论文解读

引言

​ 作者描述了当前GAN网络在医学影像数据补缺中的应用，但是目前的方法都是将数据补缺（ISM）与其他分类任务（MDM） 等分开来进行，而作者提出了一种collaborative diagnosis-synthesis framework（CDSF）即联合诊断合成框架，将两个步骤（数据补缺，模型分类）合二为一，相较于传统的做法有以下好处：

1. ISM 和 MDM 都是协同学习的，因此可以更好地协作完成合成或是分类任务。
2. MDM 可以直接为 ISM 提供多模态特征一致性，而不用再使用额外的两个单模态分类器。
3. CDSF 在神经图像合成和神经退行性疾病诊断（AD）方面实现了显着改进的性能。
   

方法

结构：

主要分为ISM和MDM模块，一个用来补充图像，一个用来分类。

**ISM：**由编码器，传输模块（RNB），解码器三部分组成。编码部分包含三个卷积 (Conv) 层，分别具有 8、16 和 32 个通道。传输部分包含六个残差网络块 (RNB)。解码部分包含两个分别具有 32 和 16 个通道的反卷积 (Deconv) 层和一个具有 1 个通道的 Conv 层。除了使用“tanh”激活的最后一个 Conv 层外，所有 Conv/Deconv 层都使用“relu”激活。这两个 Deconv 层，第二层和第三个 Conv 层以 2 的步幅构建金字塔结构。同时，对所有层进行归一化。在训练阶段，特征一致约束和体素一致约束都用于约束 ISM 的学习(参见方程(3)(4))。

**MDM:**分为特征提取部分和空间表示部分。特征提取部分有五个 Conv 层（随后是归一化和“relu”激活），分别具有 16、32、64、64 和 64 个通道。前四个 Conv 层和最后一个 Conv 层之后分别是 3×3×3 最大池化层和平均池化层，步幅为 2。特征提取部分的输出首先沿通道维度进行 l2 归一化，然后重新整形为空间表示（个人认为是按照通道进行展开，及原本输出是64*8*8,重新整形成64*64）。然后，空间表示进行 l2 归一化并馈送到具有“softmax”激活的全连接层，以计算对象属于正类和负类的概率。

**模型的总体目标：**此次生成器的损失函数来自判别器（前半部分）及对应的影像L1loss（后半部分）。而判别器的损失是四种组合的分类结果之和。

实验结果

代码复现

​ 这篇文章的代码好像并没有开源，所以个人根据自己的理解进行了复现。

​ 网络框架图：

​

ISM:

import torch
import torch.nn as nn


class rnb(nn.Module):
    def __init__(self, in_channel, out_channel, stride=1):
        super(rnb, self).__init__()
        self.conv1 = nn.Conv3d(in_channels=in_channel, out_channels=out_channel, kernel_size=3, stride=stride,
                               padding=1,
                               bias=False)
        self.bn1 = nn.BatchNorm3d(out_channel)
        self.relu = nn.ReLU()
        self.conv2 = nn.Conv3d(in_channels=out_channel, out_channels=out_channel, kernel_size=3, stride=stride,
                               padding=1, bias=False)
        self.bn2 = nn.BatchNorm3d(out_channel)

    def forward(self, x):
        identity = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        out += identity
        out = self.relu(out)

        return out


class ism(nn.Module):
    def __init__(self, in_channel=1, out_channel=1, rnb_nums=6):
        super(ism, self).__init__()
        self.in_channel = in_channel
        self.out_channel = out_channel
        self.rnb_nums = rnb_nums
        self.channel = [8, 16, 32, 32, 16, 1]

        self.encoder = nn.Sequential(
            nn.Conv3d(in_channels=self.in_channel, out_channels=self.channel[0], kernel_size=7, stride=1, padding=3),
            nn.BatchNorm3d(self.channel[0]),
            nn.ReLU(),
            nn.Conv3d(in_channels=self.channel[0], out_channels=self.channel[1], kernel_size=3, stride=2, padding=1),
            nn.BatchNorm3d(self.channel[1]),
            nn.ReLU(),
            nn.Conv3d(in_channels=self.channel[1], out_channels=self.channel[2], kernel_size=3, stride=2, padding=1),
            nn.BatchNorm3d(self.channel[2]),
            nn.ReLU()
        )

        self.rnb_layers = []
        for _ in range(self.rnb_nums):
            self.rnb_layers.append(rnb(in_channel=self.channel[2], out_channel=self.channel[3]))
        self.rnb_layers = nn.Sequential(*self.rnb_layers)

        self.decoder = nn.Sequential(
            nn.ConvTranspose3d(in_channels=self.channel[3], out_channels=self.channel[3], kernel_size=3, stride=2,
                               padding=1,output_padding=1),
            nn.BatchNorm3d(self.channel[3]),
            nn.ReLU(),
            nn.ConvTranspose3d(in_channels=self.channel[3], out_channels=self.channel[4], kernel_size=3, stride=2,
                               padding=1,output_padding=1),
            nn.BatchNorm3d(self.channel[4]),
            nn.ReLU(),
            nn.Conv3d(in_channels=self.channel[4], out_channels=self.out_channel, kernel_size=7, stride=1, padding=3),
            nn.BatchNorm3d(self.out_channel),
            nn.Tanh()
        )

    def forward(self, x):
        out = self.encoder(x)
        out = self.rnb_layers(out)
        out = self.decoder(out)
        return out


# i = torch.zeros((1, 1, 64, 64, 64))
# net = ism(1, 1, 6)
# o = net(i)
# print(o.shape)

MDM:

import torch
import torch.nn as nn
import torch.nn.functional as F


class mdm(nn.Module):
    def __init__(self, in_channal=2, num_classes=3):
        super(mdm, self).__init__()
        self.in_channal = in_channal
        self.num_classes = num_classes
        self.channels = [16, 32, 64, 64, 64, 64]

        # 2 16
        self.conv1 = nn.Conv3d(in_channels=self.in_channal, out_channels=self.channels[0], kernel_size=3, stride=1,
                               padding=1)
        self.bn1 = nn.BatchNorm3d(self.channels[0])
        self.maxpool1 = nn.MaxPool3d(kernel_size=3, stride=2, padding=1)
        # 16 32
        self.conv2 = nn.Conv3d(in_channels=self.channels[0], out_channels=self.channels[1], kernel_size=3, stride=1,
                               padding=1)
        self.bn2 = nn.BatchNorm3d(self.channels[1])
        self.maxpool2 = nn.MaxPool3d(kernel_size=3, stride=2, padding=1)
        # 32 64
        self.conv3 = nn.Conv3d(in_channels=self.channels[1], out_channels=self.channels[2], kernel_size=3, stride=1,
                               padding=1)
        self.bn3 = nn.BatchNorm3d(self.channels[2])
        self.maxpool3 = nn.MaxPool3d(kernel_size=3, stride=2, padding=1)
        # 64 64
        self.conv4 = nn.Conv3d(in_channels=self.channels[2], out_channels=self.channels[3], kernel_size=3, stride=1,
                               padding=1)
        self.bn4 = nn.BatchNorm3d(self.channels[3])
        self.maxpool4 = nn.MaxPool3d(kernel_size=3, stride=2, padding=1)
        # 64 64
        self.conv5 = nn.Conv3d(in_channels=self.channels[4], out_channels=self.channels[5], kernel_size=3, stride=1,
                               padding=1)
        self.bn5 = nn.BatchNorm3d(self.channels[5])
        self.avgpool = nn.AvgPool3d(kernel_size=3, stride=2, padding=1)

        self.relu = nn.ReLU()

        self.classifier = nn.Linear(4096, num_classes)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.maxpool1(out)
        out = self.conv2(out)
        out = self.bn2(out)
        out = self.maxpool2(out)
        out = self.conv3(out)
        out = self.bn3(out)
        out = self.maxpool3(out)
        out = self.conv4(out)
        out = self.bn4(out)
        out = self.maxpool4(out)
        out = self.conv5(out)
        out = self.bn5(out)
        out = self.avgpool(out)
        out = F.normalize(out, p=2, dim=2)
        out = torch.flatten(out, start_dim=2)
        out = F.normalize(out, p=2, dim=2)
        out = torch.flatten(out, start_dim=1)
        out = self.classifier(out)
        out = self.softmax(out)

        return out


# i = torch.zeros((1, 2, 128, 128, 128))
# net = mdm()
# print(net)
# o = net(i)
# print(o)
# print(o.shape)