医学图像分割经典模型U-Net架构全解析
U形架构由特定的编码器-解码器方案组成：编码器在每一层减少空间维度并增加通道数。另一方面，解码器增加空间维度，同时减少通道数。传递到解码器的张量通常称为“瓶颈”。最后，空间维度被恢复，以对输入图像中的每个像素进行预测。这类模型在现实世界应用中得到了广泛使用。
本文旨在探讨那些经受住时间考验的U-Net架构。
全卷积网络（FCN）
全卷积网络（FCN）是最早没有全连接层的架构之一。除了可以进行端到端训练之外，对于单个像素预测（如语义分割），它可以处理任意大小的输入。这是一种通用架构，有效地利用转置卷积作为可训练的上采样方法。
给定一个预训练的编码器，一个FCN的结构如下：
import torch
import torch.nn as nnclass FCN32s(nn.Module):def __init__(self, pretrained_net, n_class):super().__init__()self.n_class = n_classself.pretrained_net = pretrained_netself.relu = nn.ReLU(inplace=True)self.deconv1 = nn.ConvTranspose2d(512, 512, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)self.bn1 = nn.BatchNorm2d(512)self.deconv2 = nn.ConvTranspose2d(512, 256, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)self.bn2 = nn.BatchNorm2d(256)self.deconv3 = nn.ConvTranspose2d(256, 128, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)self.bn3 = nn.BatchNorm2d(128)self.deconv4 = nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)self.bn4 = nn.BatchNorm2d(64)self.deconv5 = nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)self.bn5 = nn.BatchNorm2d(32)self.classifier = nn.Conv2d(32, n_class, kernel_size=1)def forward(self, x):output = self.pretrained_net(x)x5 = output['x5'] # size=(N, 512, x.H/32, x.W/32)score = self.bn1(self.relu(self.deconv1(x5))) # size=(N, 512, x.H/16, x.W/16)score = self.bn2(self.relu(self.deconv2(score))) # size=(N, 256, x.H/8, x.W/8)score = self.bn3(self.relu(self.deconv3(score))) # size=(N, 128, x.H/4, x.W/4)score = self.bn4(self.relu(self.deconv4(score))) # size=(N, 64, x.H/2, x.W/2)score = self.bn5(self.relu(self.deconv5(score))) # size=(N, 32, x.H, x.W)score = self.classifier(score) # size=(N, n_class, x.H/1, x.W/1)return score
甚至可以从PyTorch Hub加载预训练模型：
import torch
model = torch.hub.load('pytorch/vision:v0.9.0', 'fcn_resnet101', pretrained=True)
model.eval()
注意，所有预训练模型期望输入图像以相同的方式进行归一化，即形状为(N, 3, H, W)的3通道RGB图像的小批量，其中N是图像数量，H和W预期至少为224像素。图像需要加载到[0, 1]范围内，然后使用均值mean = [0.485, 0.456, 0.406]和标准差std = [0.229, 0.224, 0.225]进行归一化。
U-Net 与 3D U-Net
随后，U-Net对FCN进行了修改和扩展。
其主要思想是让FCN在解码器的早期层中保留高级特征。为此，他们引入了长跳跃连接来定位分割。通过这种方式，编码器路径中的高分辨率特征（但语义较低）与上采样的输出相结合并重用。U-Net也是一种对称架构，如下图所示。
它可以分为编码器-解码器路径，或者等价地称为收缩-扩张路径。
编码器（左侧）：由两个3x3卷积的重复应用组成。每个卷积后面跟着一个ReLU激活函数和批量归一化。然后应用一个2x2的最大池化操作以减少空间维度。同样，在每个下采样步骤中，特征通道数加倍，而空间维度减半。
解码器路径（右侧）：扩张路径中的每一步包括特征图的上采样，然后是2x2转置卷积（将特征通道数减半）。还包括与收缩路径中对应特征图的拼接，通常还有一个3x3卷积（每个后面跟着ReLU）。在最后一层，使用1x1卷积将通道数映射到所需的类别数量。
以下是2D U-Net的实现：
import torch
import torch.nn as nn
import torch.nn.functional as Fclass DoubleConv(nn.Module):def __init__(self, in_ch, out_ch):super(DoubleConv, self).__init__()self.conv = nn.Sequential(nn.Conv2d(in_ch, out_ch, 3, padding=1),nn.BatchNorm2d(out_ch),nn.ReLU(inplace=True),nn.Conv2d(out_ch, out_ch, 3, padding=1),nn.BatchNorm2d(out_ch),nn.ReLU(inplace=True))def forward(self, x):x = self.conv(x)return xclass InConv(nn.Module):def __init__(self, in_ch, out_ch):super(InConv, self).__init__()self.conv = DoubleConv(in_ch, out_ch)def forward(self, x):x = self.conv(x)return xclass Down(nn.Module):def __init__(self, in_ch, out_ch):super(Down, self).__init__()self.mpconv = nn.Sequential(nn.MaxPool2d(2),DoubleConv(in_ch, out_ch))def forward(self, x):x = self.mpconv(x)return xclass Up(nn.Module):def __init__(self, in_ch, out_ch, bilinear=True):super(Up, self).__init__()if bilinear:self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)else:self.up = nn.ConvTranspose2d(in_ch // 2, in_ch // 2, 2, stride=2)self.conv = DoubleConv(in_ch, out_ch)def forward(self, x1, x2):x1 = self.up(x1)diffY = x2.size()[2] - x1.size()[2]diffX = x2.size()[3] - x1.size()[3]x1 = F.pad(x1, (diffX // 2, diffX - diffX // 2,diffY // 2, diffY - diffY // 2))x = torch.cat([x2, x1], dim=1)x = self.conv(x)return xclass OutConv(nn.Module):def __init__(self, in_ch, out_ch):super(OutConv, self).__init__()self.conv = nn.Conv2d(in_ch, out_ch, 1)def forward(self, x):x = self.conv(x)return xclass Unet(nn.Module):def __init__(self, in_channels, classes):super(Unet, self).__init__()self.n_channels = in_channelsself.n_classes = classesself.inc = InConv(in_channels, 64)self.down1 = Down(64, 128)self.down2 = Down(128, 256)self.down3 = Down(256, 512)self.down4 = Down(512, 512)self.up1 = Up(1024, 256)self.up2 = Up(512, 128)self.up3 = Up(256, 64)self.up4 = Up(128, 64)self.outc = OutConv(64, classes)def forward(self, x):x1 = self.inc(x)x2 = self.down1(x1)x3 = self.down2(x2)x4 = self.down3(x3)x5 = self.down4(x4)x = self.up1(x5, x4)x = self.up2(x, x3)x = self.up3(x, x2)x = self.up4(x, x1)x = self.outc(x)return x
这种方法在2D生物医学图像分割中取得了巨大成功，并且仍然被用作基线方法。那么3D图像呢？
3D-Unet
3D U-Net是在U-Net之后不久引入的，用于处理体积数据。官方示意图中只显示了3层，但在实际实现中，我们会使用更多层。每个块在卷积后使用批量归一化。
V-Net (2016)
Vnet扩展了U-Net，用于处理3D MRI体积数据。与逐片处理输入3D体积不同，该模型提出使用3D卷积。毕竟，医学图像具有固有的3D结构，逐片处理是次优的。Vnet的主要修改包括：
- 受图像分类类似工作的启发，将最大池化操作替换为步长卷积。这是通过使用步长为2的2×2×2核进行卷积来实现的。
- 在每个阶段使用5×5×5核进行带填充的3D卷积。
- 在网络的两个部分都使用了短残差连接。
- 使用3D转置卷积来增加输入的尺寸，随后是一到三个卷积层。特征图在每个解码器层减半。
以上内容可以在此图中说明：
最后，在这项工作中，引入了Dice损失，这是分割中常见的损失函数。可以在我们的开源库中找到Vnet的实现。
UNet++ (2018)
动机：U-Net中使用的跳跃连接直接将编码器的高分辨率特征图快速前馈到解码器网络。这导致了语义上不同的特征图的拼接。
UNet++背后的主要思想是在拼接之前，弥合编码器和解码器特征图之间的语义差距。为此，UNet++基于嵌套和密集的跳跃连接。UNet++可以有效地捕捉2D图像的细粒度细节。可视化如下：
UNet++由一个编码器和一个解码器组成，通过一系列嵌套的密集卷积块连接。上图由UNet++论文提供。
在上图中，黑色表示原始的U-Net，而绿色和蓝色显示跳跃路径上的密集卷积块。红色表示深度监督，意味着存在多个损失项，与标准的U-Net不同。该实现是公开可用的。
换句话说，密集卷积块使编码器特征图的语义水平更接近于解码器中等待的特征图。然而，参数数量以及训练网络的时间显著增加。这是没有类似3D卷积架构的主要原因。
No New-Net (2018)
语义图像分割中公认的U-Net基线。它在BRATS数据集上进行了测试，结果名列前茅。要点如下：
- 使用128x128x128的子体积，批大小为2
- 第一个卷积层有30个通道
- 解码器中使用三线性上采样
- 将Dice损失与负对数似然相结合
- 数据增强策略：随机旋转、随机缩放、随机弹性变形、伽马校正增强和镜像。
- L2权重衰减为10^
详细的Github仓库是可用的。
使用自动编码器正则化进行3D MRI脑肿瘤分割
尽管这不完全是传统的U-Net架构，但它值得列入此列表。编码器是一个3D ResNet模型，解码器使用转置卷积。第一个关键部分是绿色的构建块，如图所示：
它使用连续的填充3D卷积，带有组归一化、ReLU激活函数和残差跳跃连接。
第二个重要组成部分是右下角的子网络：它是一个变分自编码器，试图重建原始的3D输入图像。
为什么？
好问题！
使用自编码器分支的动机是为编码器部分提供额外的指导和正则化，因为训练数据有限。尽管该模型使用了大量的GPU内存，但它赢得了BRATS竞赛。该实现在MedicalZoo库中公开可用。
MultiResUNet：重新思考用于多模态生物医学图像分割的U-Net架构 (2020)
医学图像源自各种模态，我们关注的分割对象具有不规则和不同的尺度。为了解决这个问题，该模型提出使用类似Inception的卷积模块。以下是Inception模块工作原理的快速回顾：
遵循Inception网络的思想，他们通过将3x3和7x7卷积操作与现有的3x3卷积并行结合，来增强U-Net的多分辨率能力。
为了处理额外的网络复杂性，他们将5x5和7x7卷积层分解，使用一系列小的3x3卷积块。然后，来自三个卷积块的输出被拼接起来，以从不同尺度提取空间特征。
此外，他们逐渐增加后续卷积中的滤波器数量，以减少早期层的内存占用。最后，他们添加了一个短残差跳跃连接，并引入了一个逐点（1x1）卷积层，这可能有助于捕获额外的空间信息。
为了解决编码器-解码器特征之间的差异（由于长跳跃连接），他们提出沿着快捷连接加入一些卷积层。假设是，额外的非线性变换应该补偿解码器阶段所做的进一步处理。
这些类似Inception的架构改进在许多医学图像分割数据集中展示了优越的结果。可以自行尝试以了解更多信息。
3D U^2-Net：引入通道可分离卷积
深度可分离意味着计算跨不同通道进行。
在可分离卷积中，计算被分解为两个连续的步骤：一个通道可分离步骤独立处理每个通道，另一个1x1x通道卷积合并独立生成的特征图。
同样，通道可分离卷积对每个输入通道应用独立的卷积滤波器，如图所示：
逐点（1x1xk或1x1x1xk核）卷积线性组合每个空间位置所有通道的输出。
对于3D情况，参数增益是巨大的，尤其是在不同域（用T表示）上训练模型的多个实例时。一个具有输入通道c_in和输出通道c_out的3x3x3卷积层需要27 * c_in * c_out * T个参数。使用可分离卷积，上述滤波器中的总权重数为(27 * c_in * T) + (c_in * c_out)。
主要的假设是每个域都有自己的通道可分离滤波器，而逐点卷积核是共享的。
输入层使用16个滤波器。编码器和解码器路径都包含五个不同分辨率的层级。在每个层级内应用残差跳跃连接。跳跃连接用于从编码器对应部分保留更多上下文信息，供解码器路径使用。
显然，所提出的3D U^2-Net需要的参数最少，表明它可以跨多个域有效执行。通用模型的参数总数大约是所有独立模型的1%，而两者获得的分割精度相当。代码也是公开可用的。
结论
总而言之，没有一种模型能适合所有情况。本文尝试提供一组经过实验验证的通用想法，以便围绕U-Net展开工作。欢迎尝试其中的一些想法。有一个Git仓库收集了带有代码链接的U-Net架构。最后但同样重要的是，欢迎反馈！可以了解在我们的社交媒体页面上的看法。
参考文献
[1] Long, J., Shelhamer, E., & Darrell, T. (2015). Fully convolutional networks for semantic segmentation. In Proceedings of the IEEE conference on computer vision and pattern recognition (pp. 3431-3440).
[2] Ronneberger, O., Fischer, P., & Brox, T. (2015, October). U-net: Convolutional networks for biomedical image segmentation. In International Conference on Medical image computing and computer-assisted intervention (pp. 234-241). Springer, Cham.
[3] Milletari, F., Navab, N., & Ahmadi, S. A. (2016, October). V-net: Fully convolutional neural networks for volumetric medical image segmentation. In 2016 fourth international conference on 3D vision (3DV) (pp. 565-571). IEEE.
[4] Çiçek, Ö., Abdulkadir, A., Lienkamp, S. S., Brox, T., & Ronneberger, O. (2016, October). 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.
[5] Zhou, Z., Siddiquee, M. M. R., Tajbakhsh, N., & Liang, J. (2018). Unet++: A nested u-net architecture for medical image segmentation. In Deep Learning in Medical Image Analysis and Multimodal Learning for Clinical Decision Support (pp. 3-11). Springer, Cham.
[6] Wang, W., Yu, K., Hugonot, J., Fua, P., & Salzmann, M. (2019). Recurrent U-Net for resource-constrained segmentation. In Proceedings of the IEEE International Conference on Computer Vision (pp. 2142-2151).
[7] Huang, C., Han, H., Yao, Q., Zhu, S., & Zhou, S. K. (2019, October). 3D U^2-Net: A 3D Universal U-Net for Multi-domain Medical Image Segmentation. In International Conference on Medical Image Computing and Computer-Assisted Intervention (pp. 291-299). Springer, Cham.
[8] Ibtehaz, N., & Rahman, M. S. (2020). MultiResUNet: Rethinking the U-Net architecture for multimodal biomedical image segmentation. Neural Networks, 121, 74-87.
[9] Isensee, F., Kickingereder, P., Wick, W., Bendszus, M., & Maier-Hein, K. H. (2018, September). No new-net. In International MICCAI Brainlesion Workshop (pp. 234-244). Springer, Cham.
[10] Oktay, O., Schlemper, J., Folgoc, L. L., Lee, M., Heinrich, M., Misawa, K., ... & Glocker, B. (2018). Attention u-net: Learning where to look for the pancreas. arXiv preprint arXiv:1804.03999.
[11] Alom, M. Z., Hasan, M., Yakopcic, C., Taha, T. M., & Asari, V. K. (2018). Recurrent residual convolutional neural network based on u-net (r2u-net) for medical image segmentation. arXiv preprint arXiv:1802.06955.
[12] Myronenko, A. (2018, September). 3D MRI brain tumor segmentation using autoencoder regularization. In International MICCAI Brainlesion Workshop (pp. 311-320). Springer, Cham.
更多精彩内容 请关注我的个人公众号 公众号（办公AI智能小助手）或者 我的个人博客 https://blog.qife122.com/
对网络安全、黑客技术感兴趣的朋友可以关注我的安全公众号（网络安全技术点滴分享）
公众号二维码
公众号二维码
