在图像分割任务特别是医学图像分割中,U-Net[1]无疑是最成功的方法之一,该方法在2015年MICCAI会议上提出,目前已达到四千多次引用。其采用的编码器(下采样)-解码器(上采样)结构和跳跃连接是一种非常经典的设计方法。目前已有许多新的卷积神经网络设计方式,但很多仍延续了U-Net的核心思想,加入了新的模块或者融入其他设计理念。本文对U-Net及其几种改进版做一个介绍。
U-Net最初是一个用于二维图像分割的卷积神经网络,分别赢得了ISBI 2015细胞追踪挑战赛和龋齿检测挑战赛的冠军[2]。U-Net的一个Karas实现代码:
https://github.com/zhixuhao/unet
U-Net的结构如下图所示,左侧可视为一个编码器,右侧可视为一个解码器。编码器有四个子模块,每个子模块包含两个卷积层,每个子模块之后有一个通过max pool实现的下采样层。输入图像的分辨率是572x572, 第1-5个模块的分辨率分别是572x572, 284x284, 140x140, 68x68和32x32。由于卷积使用的是valid模式,故这里后一个子模块的分辨率等于(前一个子模块的分辨率-4)/2。解码器包含四个子模块,分辨率通过上采样操作依次上升,直到与输入图像的分辨率一致(由于卷积使用的是valid模式,实际输出比输入图像小一些)。该网络还使用了跳跃连接,将上采样结果与编码器中具有相同分辨率的子模块的输出进行连接,作为解码器中下一个子模块的输入。
3D U-Net[3]是U-Net的一个简单扩展,应用于三维图像分割,结构如下图所示。相比于U-Net,该网络仅用了三次下采样操作,在每个卷积层后使用了batch normalization,但3D U-Net和U-Net均没有使用dropout。
在2018年MICCAI脑肿瘤分割挑战赛(brats)中[4],德国癌症研究中心的团队使用3D U-Net,仅做了少量的改动,取得了该挑战赛第二名的成绩,发现相比于许多新的网络,3D U-Net仍然十分具有优势[5]。3D U-Net的一种Pytorch实现:
https://github.com/wolny/pytorch-3dunet
TernausNet全称为"TernausNet: U-Net with VGG11 Encoder Pre-Trained on ImageNet for Image Segmentation"[6]。该网络将U-Net中的编码器替换为VGG11,并在ImageNet上进行预训练,从735个参赛队伍中脱颖而出,取得了Kaggle 二手车分割挑战赛(Carvana Image Masking Challenge)第一名。代码链接:
https://github.com/ternaus/TernausNet
下图是该网络的示意图:
Res-UNet和Dense-UNet分别受到残差连接和密集连接的启发,将UNet的每一个子模块分别替换为具有残差连接和密集连接的形式。[6] 中将Res-UNet用于视网膜图像的分割,其结构如下图所示,其中灰色实线表示各个模块中添加的残差连接。
密集连接即将子模块中某一层的输出分别作为后续若干层的输入的一部分,某一层的输入则来自前面若干层的输出的组合。下图是[7]中的密集连接的一个例子。该文章中将U-Net的各个子模块替换为这样的密集连接模块,提出Fully Dense UNet 用于去除图像中的伪影。
MultiResUNet[8]提出了一个MutiRes模块与UNet结合。MutiRes模块如下图所示,是一个残差连接的扩展,在该模块中三个3x3的卷积结果拼接起来作为一个组合的特征图,再与输入特征图经过1x1卷积得到的结果相加。
该网络的结构图如下图所示,其中各个MultiRes模块的内部即为上图所示。
该网络除了MultiRes模块以外,还提出了一个残差路径(ResPath), 使编码器的特征在与解码器中的对应特征拼接之前,先进行了一些额外的卷积操作,如下图所示。作者认为编码器中的特征由于卷积层数较浅,是低层次的特征,而解码器中对应的特征由于卷积层更深,是较高层次的特征,二者在语义上有较大差距,推测不宜直接将二者进行拼接。因此,使用额外的ResPath使二者在拼接前具有一致的深度,在ResPath1, 2, 3, 4中分别使用4,3,2,1个卷积层。
该文章在ISIC、CVC-ClinicDB、Brats等多个数据集上验证了其性能。代码链接为
https://github.com/nibtehaz/MultiResUNet
R2U-Net全称叫做Recurrent Residual CNN-based U-Net[9]。该方法将残差连接和循环卷积结合起来,用于替换U-Net中原来的子模块,如下图所示
其中环形箭头表示循环连接。下图表示了几种不同的子模块内部结构图,(a)是常规的U-Net中使用的方法,(b)是在(a)的基础上循环使用包含激活函数的卷积层,(c)是使用残差连接的方式,(d)是该文章提出的结合(b)和(c)的循环残差卷积模块。
该方法也在皮肤病图像、视网膜图像、肺部图像等几个公共数据集验证了其性能,代码链接:
https://github.com/LeeJunHyun/Image\_Segmentation#r2u-net
Attention UNet[10]在UNet中引入注意力机制,在对编码器每个分辨率上的特征与解码器中对应特征进行拼接之前,使用了一个注意力模块,重新调整了编码器的输出特征。该模块生成一个门控信号,用来控制不同空间位置处特征的重要性,如下图中红色圆圈所示。
该方法的注意力模块内部如下图所示,该模块通过1x1x1的卷积分别与ReLU和Sigmoid结合,生成一个权重图, 通过与编码器中的特征相乘来对其进行校正。
下图展示了注意力权重图的可视化效果。从左至右分别是一幅图像和随着训练次数的增加该图像中得到的注意力权重。可见得到的注意力权重倾向于在目标器官区域取得大的值,在背景区域取得较小的值,有助于提高图像分割的精度。
该文章的代码链接:
https://github.com/ozan-oktay/Attention-Gated-Networks
基于U-Net框架设计的图像分割网络还有很多,难以一一列举,这里再提供两篇具有参考性的文章:
AnatomyNet: Deep 3D Squeeze-and-excitation U-Nets for fast and fully automated whole-volume anatomical segmentation
H-DenseUNet: Hybrid Densely Connected UNet for Liver and Liver Tumor Segmentation from CT Volumes
UNet是一个经典的网络设计方式,在图像分割任务中具有大量的应用。也有许多新的方法在此基础上进行改进,融合更加新的网络设计理念,但目前几乎没有人对这些改进版本做过比较综合的比较。由于同一个网络结构可能在不同的数据集上表现出不一样的性能,在具体的任务场景中还是要结合数据集来选择合适的网络。希望本文对做图像分割的同学有所启发。
[1] Ronneberger, Olaf, Philipp Fischer, and Thomas Brox. "U-net: Convolutional networks for biomedical image segmentation." In International Conference on Medical image computing and computer-assisted intervention, pp. 234-241. Springer, Cham, 2015.
[2] U-Net: Convolutional Networks for Biomedical Image Segmentation
(链接:https://lmb.informatik.uni-freiburg.de/people/ronneber/u-net/ )
[3] Çiçek, Özgün, Ahmed Abdulkadir, Soeren S. Lienkamp, Thomas Brox, and Olaf Ronneberger. "3D U-Net: learning dense volumetric segmentation from sparse annotation." In International conference on medical image computing and computer-assisted intervention, pp. 424-432. Springer, Cham, 2016.
[4]Section for Biomedical Image Analysis (SBIA)
( 链接:https://www.med.upenn.edu/sbia/brats2018/data.html )
[5] Isensee, Fabian, Philipp Kickingereder, Wolfgang Wick, Martin Bendszus, and Klaus H. Maier-Hein. "No new-net." In International MICCAI Brainlesion Workshop, pp. 234-244. Springer, Cham, 2018.
[6] Xiao, Xiao, Shen Lian, Zhiming Luo, and Shaozi Li. "Weighted Res-UNet for High-Quality Retina Vessel Segmentation." In 2018 9th International Conference on Information Technology in Medicine and Education (ITME), pp. 327-331. IEEE, 2018.
[7] Guan, Steven, Amir Khan, Siddhartha Sikdar, and Parag V. Chitnis. "Fully Dense UNet for 2D Sparse Photoacoustic Tomography Artifact Removal." arXiv preprint arXiv:1808.10848 (2018).
[8] Ibtehaz, Nabil, and M. Sohel Rahman. "MultiResUNet: Rethinking the U-Net Architecture for Multimodal Biomedical Image Segmentation." arXiv preprint arXiv:1902.04049 (2019).
[9] Alom, Md Zahangir, Mahmudul Hasan, Chris Yakopcic, Tarek M. Taha, and Vijayan K. Asari. "Recurrent residual convolutional neural network based on u-net (r2u-net) for medical image segmentation." arXiv preprint arXiv:1802.06955 (2018).
[10] Oktay, Ozan, Jo Schlemper, Loic Le Folgoc, Matthew Lee, Mattias Heinrich, Kazunari Misawa, Kensaku Mori et al. "Attention U-Net: learning where to look for the pancreas." arXiv preprint arXiv:1804.03999 (2018).