【论文笔记】ASNet：基于生成对抗网络（GAN）的无监督单模和多模配准网络（范敬凡老师）

2020-07-27

本文是论文《Adversarial learning for mono- or multi-modal registration》的阅读笔记，是范敬凡老师的工作，是文章《Adversarial Similarity Network for Evaluating Image Alignment in Deep Learning based Registration》中工作的延续。关于后一篇文章，可以查看我的博客。

文章提出了一个基于GAN（生成对抗网络）的无监督配准模型，需要任何ground-truth形变场，也不需要指定相似性度量标准。

一、相关工作

配准算法是来获取一个可以配准图像的空间变换，空间变换通常包括线性变换（平移、旋转、缩放、shearing）和非线性变换（体素到体素的关系）。配准问题可以表示为一个获取使得图像不相似性最小的形变场$\phi^ $，表示为下式：
$$
\phi^{}=\underset{\phi}{\operatorname{argmin}} \operatorname{dissim}\left(I_{M} \circ \phi, I_{F}\right)+\lambda \operatorname{reg}(\phi)
$$
其中，$I_M$和$I_F$分别表示浮动图像和固定图像；$I_M\circ\phi$表示根据形变场$\phi$变形后的浮动图像；图像的不相似性$dissim(I_M\circ\phi,I_F)$可以定义为SSD、CC/NCC、MI/NMI等；$reg(\phi)$是让形变场保持平滑的正则化项；$\lambda$是平衡权重。形变场的正则化可以通过高斯平滑，使用样条曲线或微分同胚来实现。

传统的医学图片配准方法通常是高维的数学优化，所以很耗时间。基于深度学习的有监督配准方法多聚焦于预测体素到体素的映射，即密集形变场（dense deformation fields），但是它们效果又被ground-truth形变场所限制。而无监督的配准方法旨在通过最大化图像之间的相似度来学习一个形变场，这里的图像相似度通常采用SSD（平方误差和）、CC（互相关）等，但是很难决定在不同的配准问题中采取哪种特定的相似性度量才会达到最好效果。无监督的方法另一个优点是没有有监督信息，所以节省了内存/显存，进一步同样的显存可以容纳更大的图像，使得大尺度的无标签图像的配准变成可能。

二、网络结构

该模型实现了基于patch和基于图像的单模和多模3D图像的配准，单模图像选用的是脑部MR图像，多模选用的是骨盆的MR图像和CT图像。文章的贡献如下：

与传统的配准方法相比，该方法更鲁棒，更快，并且是端到端的；
与有监督配准方法相比，该方法不需要ground-truth形变场；
该方法可以自动学习图像之间的相似性度量，而不需要任何先验假设；
可以同时处理单模和多模配准问题。

文章提出的基于GAN的无监督配准模型包括一个配准网络（生成器）、一个判别器和空间变换网络。下图是网络的整体结构示意图。

1. 配准网络

配准网络以两张图片作为输入，输出为和输入大小相同的形变场。配准网络R学习一个形变场，可以记为$R:(I_M,I_F)\rightarrow\phi$，配准网络通过最大化图像间的相似性来训练。

配准网络是基于U-Net的，在编码器阶段每个卷积的卷积核大小为$3\times3\times3$，并且后面跟着ReLU激活函数，然后使用$2\times2\times2$的最大池化层进行下采样。重复以上步骤两次就是编码器的结构。在解码器阶段，使用$2\times2\times2$的反卷积层进行上采样，每个反卷积层后也跟着一个ReLU激活函数。在编码器和解码器之间有跳跃连接。最终的形变场通过$1\times1\times1$的全卷积层（没有激活函数）来得到。

下图是配准网络的结构示意图。

下图是基于patch和基于全图像的配准中配准网络和判别器网络（没有标出）的具体结构：

2. 判别网络

判别器是基于CNN的，判别器用来判别两张图片是否已经对齐，并在训练时将没对齐的信息传递给配准网络（通过反向传播更新配准网络的参数），即图像相似性由判别网络D来决定，判别网络判别两幅图像的相似性，并得到一个相似性概率$p\in[0,1]$，1表示相似，0表示不相似。

下图是判别器网络的结构示意图，如图所示，每个卷积层是0填充的，并且跟着ReLU激活函数，然后使用两个最大池化，最后使用全连接层加sigmoid激活函数来聚合所有的信息得到最终的输出。

3. 空间变换网络

配准网络和判别器之间由空间变换网络（STN）连接，空间转换层可以根据形变场来对浮动图像做变形，使用三线性插值，该过程可表示为：
$$
\begin{aligned}
I_{W}(x) &=I_{M}(x+\phi(x)) \
& \approx \Sigma_{y \in \mathcal{N}(x+\phi(x))} I_{M}(y) \cdot \prod_{d \in{0,1,2}}\left(1-\left|x_{d}+\phi_{d}(x)-y_{d}\right|\right)
\end{aligned}
$$
其中，$I_W$是浮动图像变形后的图像，$x$是体素位置，$y \in \mathcal{N}(x+\phi(x))$是位置$x+\phi(x)$的8体素的立方体邻居，$d$是图像空间的维度。空间转换层没有可训练的参数。

三、训练

在训练时，按照以下顺序进行训练：

先喂给判别器一个参考图像（reference image）和一个固定图像（fixed image），参考图像就是和固定图像很相似的图像，来学习配准好的图像是什么样的；
喂给判别器一个根据预测的形变场变形后的浮动图像和一个固定图像，让判别器学习没有配准好的图像是什么样的；
喂给配准网络一个浮动图像和一个固定图像来学习能在判别器得到高分数的形变场。

当判别器不能区分正例样本和负例样本时，训练就收敛了。

1. 判别器的训练

判别器有两个输入样例，一个是正例$P^+$，即一对配准好的图像；另一个是负例$P^-$，即一对没有配准好的图像，判别器的损失如下：
$$
\mathcal{L}_{D}(p)=\left{\begin{array}{ll}
-\log (p), & \quad c \in P^{+} \
-\log (1-p), & c \in P^{-}
\end{array}\right.
$$
其中$p$是判别器的输出，即图像相似性概率；$c$表示输入样例，正例的值应该接近于1，而负例的值应该接近于0。

正例

对于单模配准来说，理想的正例样本是两张完全相同的图像，但是这是不现实的，所以选用的参考图像是通过浮动图像和固定图像加权得到的，即：
$$
I_{R}=\alpha \cdot I_{M}+(1-\alpha) \cdot I_{F}, 0<\alpha<1
$$
在训练的初始阶段将$\alpha$设为0.2，以弱化相似性的需要，在后期将$\alpha$设为0.1以达到更高的精度。

对应多模配准来说，参考图像和浮动图像是来自同一模态的，而固定图像来自不同模态。使用同一模态下的少量MR和CT的图像对作为参考图像。

下图中的(a)、(b)分别表示单模和多模配准中的正例和负例样本。

2. 配准网络的训练

配准网络中相似性损失定义为：
$$
\mathcal{L}{R}(p)=-\log (p), c \in P^{-}
$$
形变场的平滑性定义为：
$$
\mathcal{L}{\mathrm{reg}}(\phi)=\sum_{v \in \mathbb{R}^{3}} \nabla \phi^{2}(v)
$$
其中，$v$表示体素位置，此外，反折叠惩（anti-floding）罚项被用作正则项来惩罚较大的折叠（$\nabla\phi(v)+1<0$）。配准网络的总损失如下：
$$
\mathcal{L}=\mathcal{L}{R}(p)+\lambda \cdot \mathcal{L}{\mathrm{reg}}(\phi)
$$
其中$\lambda$被设置为1000。

四、实验

1. 具体实施

该网络实现了基于patch 的和基于图像的配准两种方式，基于patch 的配准使用有重叠的大小为$68\times68\times68$的图像patch作为输入，输出大小为$28\times28\times28$的形变场向量patch。特意只预测中心的形变场，这是因为中心的信息被利用的最好，而边界信息也有所缺失，置信度不高。在获取输入patch的时候，以步长为28取patch块，也就是和输出大小相同，这样就可以保证输出的patch可以无重叠的拼成一个完整的形变场。

2. 实验设置

单模和多模配准中使用的网络结构是相同的，使用的优化器都是Adam，学习率为$1e^{-4}$，每经过50K次迭代以0.5的权重进行衰变。基于patch的配准训练10个epoch，基于图像的训练20个epoch。

使用的baseline为LCC-demons和SyN模型，使用的相似性度量为SSD和CC。DL-GT表示使用由SyN产生的ground-truth形变场的深度学习配准模型，DL-SSD和DL-CC分别表示选用SSD和CC作为相似性度量的深度学习模型，DL-ASN是本文提出的模型。除非特别说明，所有的基于学习的方法都是基于patch的。选用的数据集有LPBA40、IBSR18、CUMC12和MGH10。

在单模配准中，在预处理阶段，对图像进行了线性配准（仿射配准），并全部重采样到$224\times224\times160$大小，训练集使用的是LPBA40的30张图像，LPBA40另外10张图像用作验证集，剩余的数据集图像全用作测试集，这样做是为了验证模型的通用性。评价指标为DSC（Dice相似性系数）和ASD（平均表面距离）。

多模配准时，使用的是22个前列腺癌患者的骨盆MR和CT图像，大小为$224\times192\times96$，使用的相似性度量为MI（互信息），这是因为SSD和CC在多模配种中不适用。在训练阶段，随机选择来自相同患者的10个图像对作为预备图像对，然后选择$15\times15$个来自CT和MR的图像对形成训练集，剩余的7个图像对作用测试集。