CV04-UNet笔记

目录

一、UNet模型

二、Encoder & Decoder

2.1 Encoder

2.2 Decoder

2.3 classifier

---

学习U-Net: Convolutional Networks for Biomedical Image Segmentation，记录一些自己认为重要的要点，以免日后遗忘。

代码：https://github.com/Ascetics/LaneSegmentation/blob/master/nets/unet.py

参考pytorch官方实现和文档做出自己的实现https://github.com/Ascetics/LaneSegmentation/blob/master/nets/unet.py

一、UNet模型

医学图像分类处理往往细化到像素级别。不仅仅是医学领域，语义分割都是像素级别的分类。于是U-Net诞生了，论文题目《U-Net：用于生物医学图像分割的卷积网络》，但是U-Net不仅仅是用于医学图像处理。

FCN的思想是特征提取阶段进行卷积、池化，图像分辨率减小，特征提取后面接上几层上采样（转置卷积）来提高分辨率。为提高分割精确度，将特征提取阶段的输出和上采样阶段输出做特征融合，这样网络就能学习到更精确的输出。

U-Net在FCN的基础上改造，需要的训练集更小，分割精确度更高。改造的部分是：上采样中仍保持特征的channel很大，特征融合直接使用特征提取阶段的输出，特征提取的输出直接concatenate到上采样输出的后面，保证更多的信息传递到分类阶段。于是形成了如论文图片中这样的U型结构。

U-Net左边是encoder阶段，由5个相似结构encode block构成。第一个encode block，2个不加padding的3×3卷积，每个卷积后面跟BN和ReLU激活函数；第二个到第五个encode block，先maxpool下采样，分辨率减小一倍通道增大一倍，再接2个不加padding的3×3卷积，每个卷积后面跟BN和ReLU激活函数。这5个encode block，除了最后一个之外，输出都要保留，再decoder里作为shortcut使用。

U-Net右边是decoder阶段，由相似结构decode block构成。每个结构里面先是转置卷积上采样2倍（kernel_size=2, stride=2），然后在后面拼接经过剪裁的对应encode block的输出（shortcut），然后是2个不加padding的3×3卷积，每个卷积后面跟BN和ReLU激活函数。每个decode block结构将通道减少一倍。

最后一个1×1卷积，将channel调整到n_class分类。

注意：图中对encoder阶段的剪裁是必要的，由于不加padding的卷积操作、池化等造成了边界像素的损失，必须经过剪裁才能让两个张量spatial大小一致，才能拼接。

二、Encoder & Decoder

Encoder & Decoder 实际上就是下采样、上采样的一种方式。U-Net提出的Encoder & Decoder并不是一定之规，encoder和decoder都可以根据需要用其他模型作为backbone进行替换，比如可以用ResNet作为backbone。记录一下Encoder & Decoder的细节，以及扩展应用方式。

2.1 Encoder

U-Net论文里面encoder的卷积操作不加padding，实际工程中可以改成加padding变成same卷积。

U-Net使用预训练的模型。U-Net更像一个框架，而不是一个模型。论文里面下采样每个结构都是两个卷积，也可以把这个结构替换成ResNet、ResNeXt、MobileNet等已经训练好的模型。Encoder中encode block的数量也不一定是5个，可以根据自己的需要加深。

以ResNet为例。

• 第1个encode block和input一致，所以直接用输入作为输出，pytorch中用空的nn.Sequencial作为encode block。输出shortcut0。
• 第2个encode block下采样2倍，所以用ResNet的Conv1、BN、ReLU作为encode block。输出shortcut1。
• 第3个encode block下采样2倍，所以用ResNet的maxpool、layer1（论文Conv2_x不做下采样）作为encode block。输出shortcut2。
• 第4个encode block下采样2倍，所以用ResNet的layer2（论文Conv3_x）作为encode block。输出shortcut3。
• 第5个encode block下采样2倍，所以用ResNet的layer3（论文Conv4_x）作为encode block。输出shortcut4。
• 第6个encode block下采样2倍，所以用ResNet的layer4（论文Conv5_x）作为encode block。输出不作为shortcut。

2.2 Decoder

U-Net的decoder是和encoder对称的逐层上采样，每个上采样后面跟了相似的结构，都是2个3×3卷积作为decode block，这就增加了decode的深度。

U-Net的上采样可以使用Transposed Convolution转置卷积，也可以使用Bilinear Interpolation双线性差值。对于不同的数据集可能要采取不同的上采样策略。

U-Net的特征融合，使用的是拼接concatenate，这与FCN的直接相加是不同的。

以ResNet为例。

• 第1个decode block。用第6个encode block输出上采样2倍，和shortcut4做cat，然后2个3x3卷积。最终输出的channels应该是第6个encode block输出channels的一半。
• 第2个decode block。用第1个decode block输出上采样2倍，和shortcut3做cat，然后2个3x3卷积。最终输出的channels应该是第1个decode block输出channels的一半。
• 第3个decode block。用第2个decode block输出上采样2倍，和shortcut2做cat，然后2个3x3卷积。最终输出的channels应该是第2个decode block输出channels的一半。
• 第4个decode block。用第3个decode block输出上采样2倍，和shortcut1做cat，然后2个3x3卷积。最终输出的channels应该是第3个decode block输出channels的一半。
• 第5个decode block。用第4个decode block输出上采样2倍，和shortcut0做cat，然后2个3x3卷积。最终输出的channels应该是第4个decode block输出channels的一半。

2.3 classifier

最后一个decode block还需要经过1个1x1卷积，将输出channels调整到n_class个分类。