【论文翻译-2】BiO-Net: Learning Recurrent Bi-directional Connections for Encoder-Decoder Architecture

BiO-Net: Learning Recurrent Bi-directional Connections for Encoder-Decoder Architecture

• 源自：MICCAI2020（Medical Image Computing and Computer Assisted ）

• 阅读日期：2020年11月23日

Abstract

• U-Net是深度学习解决计算机视觉任务的基本方法，可以解决例如语义分割、超分辨率、图像去噪、图像复原之类的任务。之前对于U-Net的改进主要是修改原有的模块或者是使用新的模块来提高性能。这些方法通常不可避免增大了模型的复杂性。为了解决这个问题，本文提出了新的双向O形状的网络（Bi-directional O-shape network，BiO-Net），用循环的方法复用构建的模块而不增加参数。本文提出的双向跳跃连接能直接引用到所有编码器解码器体系结构中。
• Code：https://github.com/tiangexiang/BiO-Net.

1 Introduction

• 近年来，基于深度学习的方法在辅助医学图像分析中很流行，例如整个切片图像的分类，脑损伤分割，医学图像合成。作为最流行的基于深度学习的模型之一，U-Net在众多医学图像计算研究中展示了其令人印象深刻的表现能力。U-Net引入了跳跃连接，它跨越多个语义尺度聚合特征表示，能够防止信息丢失。

• **U-Net的变体：**最近的工作提出以不同的模块设计和网络构造来扩展U-Net结构，这表明了它在各种视觉分析任务上的潜力。V-Net在高维体素上应用U-Net，并保留原始的内部结构。W-Net通过用一个自动编码模型连接两个U-Net用于解决无监督分割问题。与U-Net相比，M-Net在不同的层次上添加不同尺度的输入特征，通过一系列的上采样和下采样层捕捉多层次的视觉细节。最近，**U-Net++**采用嵌套密集跳跃连接，更有效地表示了细粒度对象的细节。此外，attention U-Net使用额外的分支自适应地将注意力机制应用于跳跃和解码特征的融合。然而，这些方法引入额外的构建模块，这将导致更多的网络参数，从而增加GPU内存。与上述变体不同，本文提出的BiO-Net通过一种新的特征复用机制提升了U-Net的性能，在编码器和译码器之间建立双向连接，用循环的方式来推断。

• **Recurrent Convolutional Networks：**使用循环卷积迭代地细化在不同时间提取的特征已经被证明对于许多计算机视觉问题是可行和有效的。Guo等人提出在ResNet中复用残差块，以充分利用可用参数，显著减小模型大小。这种机制也有利于U-Net的发展。因此，Wang等人提出R-U-Net，它将U-Net的多个配对编码器和解码器反复连接起来，以增强其对于语义分割的识别能力，但同时也引入了额外的可学习块。BiO-Net与R-U-Net区别在于反向跳跃连接的设计，反向跳跃连接中每个解码层的潜在特征都被复用，使更多的中间信息聚合能够在时间步间保持梯度。R2U-Net采取相似的方法，但只在细化的每个层上循环最后的模块。与之相反，我们的方法在现有的编码器和译码器中学习循环链接，在涉及精细解码特征的情况下循环相同层的模块。

• 最后，我们提出一种带有双向O形状循环U-Net的推理网络（BiO-Net），通过反向跳跃连接将解码器特征映射回编码器，并在编码器和解码器之间循环。与之前的工作相比较，我们的方法实现了更好的特征细化，在我们的网络中产生了多次编码和解码过程。将BiO-Net应用于细胞语义分割任务和EM膜语义分割任务，我们的结果显示，我们提出的模型优于其他U-Net变体，包括许多先进方法。在超分辨率任务中也证明我们的模型应用于不同场景的意义。

2 Methods

• 如图1所示，BiO-Net采用与U-Net相同的网络架构，不使用任何额外的模块，而是使用成对的双向连接。随着$t$的增加，在展开过程中不引入额外的训练参数，性能得到了提高。此外，我们的方法不局限于U-Net，可以集成到其他的编码器解码器架构以完成各种视觉分析任务。

2.1 Recurrent Bi-directional Skip Connections

• 我们的BiO-Net模型独特之处在于引入了双向跳跃连接，便于编码器处理解码器中的语义特征，反之亦然。

• **Forward Skip Connections：**向前跳跃连接在相同层连接编码器和解码器，能够保留编码器的低级视觉特征$f_{enc}$和保留梯度。因此，第$l$个解码器模块能够融合$f_{enc}$和其来自低层模块的输入${\hat{x}}_{in}$，并通过解码卷积$DEC$传播从而产生$f_{dec}$，通过$UP$模块恢复到高分辨率。该过程定义为：
  $$f_{dec}=DEC([f_{enc},{\hat{x}}_{in}])\text{\hspace{1em}(1)}$$
  其中，使用级联作为融合机制$[\cdot ]$，为了简化，本文省略了编码器和译码器的索引符号$l-th$，即第$l$层

• **Backward Skip Connections：**在所提出新的反向跳跃连接的帮助下，它将解码后的高级语义特征$f_{dec}$从解码器传递到编码器，我们的编码器能够结合$f_{dec}$和它原来由之前块产生的输入$x_{in}$，从而实现低级视觉特征和高级语义特征的聚合。类似于向前跳跃连接的解码路径，编码过程能够使用编码卷积$ENC$来定义：
  $$f_{enc}=ENC([f_{dec},x_{in}])\text{\hspace{1em}(2)}$$
  其中$DOWN$模块将$f_{enc}$提供给后续的编码器，以便进行更深入的特征提取

• **Recursive Inferences：**上述的双向跳跃连接为编码器解码器架构创建一个O形推理路径，值得注意的是，这个O形推理路径能够被循环多次，更重要的是循环传播策略不会引入额外的训练参数。因此本文提出的O-形连接的编码器和解码器的当前第$i$个循环的输出表示如下：
  $$\begin{gathered} x_{out}^i=DOWN(ENC([DEC([f_{enc}^{i-1},{\hat{x}}_{in}^{i-1}]),x_{in}^i])), \\ {\hat{x}}_{out}^i=UP(DEC([ENC([f_{dec}^i,x_{in}^i]),{\hat{x}}_{in}^i])).\text{\hspace{1em}(3)} \end{gathered}$$
  与原始的U-Net相比，BiO-Net同时考虑了编码和解码的特征，并根据前一次循环的特征进行细化

2.2 BiO-Net Architecture

• 我们在网络架构种使用简单的卷积层，批量归一化层，ReLU层。除了批量归一化层，其他层都被复用。首先将输入的图像送入三个卷积块序列中提取低级特征。注意到，第一个模块中没有连接反向跳跃连接，因此第一个模块的参数不会在循环的时候被复用。接下来将提取的特征用最大池化进行特征下采样。在编码阶段，通过成对的正向和反向连接重复使用参数和块，如图1所示。在编码阶段结束后，中间阶段包含卷积块，用于进一步细化编码特征。然后，特征被传得到一系列解码块中，这些解码块使用反卷积操作来恢复已编码的细节。在解码阶段，我们的反向跳跃连接保存了提取的特征，通过将这些提取的特征与编码器相同层的特征连接起来，如2.1节所示。循环从解码阶段最后一个卷积块生成的输出开始。在循环编码和解码阶段之后，更新的输出将送入最后一个模块，类似的，最后一个模块不参与到循环中。

• 为了更好得表示，我们定义了几何符号：$t$表示总的循环时间；$\times n$表示对所有隐藏层输出的通道数量扩大$n$倍；$w$表示从最深的编码层开始算起，使用的反向跳跃连接的数量；$INT$表示集合每次循环产生的解码特征，将他们作为一个整体送入最后一个模块；$l$表示最大编码深度。图2是示意图。

3 Experiments

• Datasets：本文提出的方法在三种常见的数字病理图像分析任务进行评估：细胞核分割、EM膜分割、图像超分辨率，总共四个不同的数据集。①评估细胞核分割有两个公共数据集：MoNuSeg和TNBC。MoNuSeg数据集包括30个图像的训练集和14个图像的测试集，图像的大小为$1000^2$（从多个器官的不同整体切片中采样）。我们从每个图像的四个角落中提取了$512^2$的补丁，使数据集扩大了4倍。TNBC由50张大小为$512^2$的组织病理学图像组成，没有指定测试集。两个数据集都有像素级别的标注。②评估EM膜分割的数据集：小鼠梨状皮层EM数据集，包括四个EM图像堆，图像切片大小分别是$255^2$，$512^2$，$512^2$和$256^2$。③评估图像超分辨率的数据集：MICCAI15 CBTC。我们在40倍放大水平下，采样了大小为$512^2$的2900个补丁，训练集和测试集的比例为9:1。

• **Implementation Details：**我们使用Adam优化器，初始学习率为0.01，衰减率为0.00003，从而最小化交叉熵损失（分割任务）或均方差（超分辨率任务）。训练集的数据增强使用了随机旋转（[-15°, +15°]），随机移动（x和y方向范围为[-5%, +5%]），随机裁剪，随机缩放（范围[0, 0.2]），随机翻转（水平和垂直）。在训练和测试阶段，batch size都设为2。除非有说明，BiO-Net默认编码深度为4，并且网络的每层都有反向跳跃连接。在所有实验中，网络训练了300个周期，使用单个NVIDIA GeForce GTX 1080 GPU，在Keras上运行。考虑到GPU的限制，我们令$t=3$为最大可能时间步长。

3.1 Semantic Segmentation

• **Nuclei Segmentation：**在这项任务中，我们的方法与基准U-Net以及其他先进方法比较。【6】之后，模型仅在MoNuSeg的训练集上训练，在MoNuSeg的测试集和TNBC数据集上进行测试。使用DICE（骰子系数）和IoU（交并集）进行评估。如表1所示，我们在MoNuSeg测试集中的结果好于其他方法。我们在TNBC数据集上的结果也比其他方法高，这表明了强大的泛化能力。此外，与其他U-Net变体相比，BiO-Net更节约内存。我们方法和同是循环网络的R2U-Net定性比较如图3所示。可以发现，随着推理时间的增加，我们的模型分割的细胞核更精确。在实验中，BiO-Net在$t=1$，$t=2$，$t=3$时产生两个batch预测结果的时间分别为35，52，70。将我们的方法结合另一个编码器解码器架构LinkNet，该评估也在表中，我们的BiO-LinkNet增加了元素级别的跳跃连接，因此与原始的LinkNet的参数数量相同。

• 表2展示燃烛实验。结果表明，采用所提出的双向跳跃连接在编码器和解码器之间进行循环处理，总体上提高了网络性能。将所有推理循环得到的解码特征集成在一起，可以在两个数据集中产生最好的性能。此外，我们发现当网络中参数不足时，增加推理循环次数效果不大，甚至会使结果变差。同样有趣的是，当用较浅的编码深度构建BiO-Net时，我们的模型在这两个数据集上的表现要比编码深度更深的模型好。

• **EM Membrane Segmentation：**我们通过分割小鼠梨状皮质EM图像进一步评估我们的方法，其中模型在stack1和stack2上进行训练，并在stack4和stack3上进行验证。结果采用Rand F-score【2】进行评价。如表3所示，我们的方法用所提出的双向o形跳跃连接分割效果更好。

3.2 Super Resolution

• 在超分辨率任务中，使用低分辨率(下采样)图像作为输入，训练网络向原始的高分辨率ground truth进行训练，可以在低分辨率的基础上恢复缺失的细节，生成高分辨率的组织病理学图像，从而帮助医学成像分析。采用两种先进的方法，FSRCNN和SRResNet，与我们的网络进行比较。定性结果以及整个测试集的峰值信噪比(PSNR)分数如图4所示。可以看出，我们的方法比目前最先进的方法产生的边缘更靠谱，这验证了我们的网络应用于不同视觉任务的可行性。

4 Conclusion

• 本文提出了一种新型的U-Net循环变体，称为BiO-Net。该网络在不引入额外训练参数的情况下，比U-Net拥有更好的性能，它利用成对的前向和反向跳跃连接来组成编码器和解码器之间复杂的关系。该模型在训练和推理中能对参数进行循环复用，在语义分割和超分辨率任务中大量的实验表明所提出模型的有效性，在不引入额外参数的条件下优于U-Net及其变体。

Reference

[16] Mehta, R., Sivaswamy, J.: M-net: A convolutional neural network for deep brain structure segmentation. In: 14th International Symposium on Biomedical Imaging (ISBI). pp. 437–440. IEEE (2017)
[17] Milletari, F., Navab, N., Ahmadi, S.A.: V-net: Fully convolutional neural networks for volumetric medical image segmentation. In: 4th International Conference on 3D Vision (3DV). pp. 565–571. IEEE (2016)
[25] Xia, X., Kulis, B.: W-net: A deep model for fully unsupervised image segmentation. arXiv preprint arXiv:1711.08506 (2017)
[28] Zhou, Z., Siddiquee, M.M.R., Tajbakhsh, N., Liang, J.: Unet++: A nested u-net architecture for medical image segmentation. In: Deep Learning in Medical Image Analysis and Multimodal Learning for Clinical Decision Support (DLMIA), pp. 3–11. Springer (2018)