【U-Net】Pytorch实现

2021年11月06日11:16:25
今天来完成U-Net

原文链接：https://arxiv.org/abs/1505.04597
论文题目：U-Net: Convolutional Networks for Biomedical Image Segmentation（2015）
作者：Olaf Ronneberger, Philipp Fischer, and Thomas Brox
代码：U-Net网络结构讲解(语义分割)
推荐阅读Pytorch-GPU安装教程大合集（Perfect完美系列）

原文插图：

网络结构：

  网络架构如上图1所示。它由收缩路径（左侧）和扩张路径（右侧）组成。收缩路径遵循卷积网络的典型结构。它包括重复应用两个3x3卷积（未填充），每个卷积后面都有一个非线性激活单元ReLU和一个2x2最大池运算，步长为2用于下采样。在每个下采样步骤中，我们将特征通道的数量增加一倍。扩展路径中的每一步包括特征图的上采样，然后是2x2卷积（反卷积），将特征通道的数量减半，与收缩路径中相应裁剪的特征图进行拼接，以及两个3x3卷积，每个卷积后是一个ReLU。由于每次卷积都会丢失边界像素，因此需要进行裁剪。在最后一层，使用1x1卷积将每个分量特征向量映射到所需数量的类。该网络总共有23个卷积层。

结构分析：
  U-net整个网络结构成一个U字型，两边对称；左半边进行卷积-激活-池化，每次池化时通道的维度翻倍；右边进行的是卷积-激活-上采样（反卷积），每次上采样的时候通道维度减半然后把左边的部分进行copy-and-crop(复制然后拼接到上采样的结果)，直到最后一层使用1*1卷积来将64维的特征图调整到2维
Pytorch：U-Net

"""
Author: yida
Time is: 2021/11/6 15:06 
this Code: 实现U-Net
"""
import os

import torch
import torch.nn as nn

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"


class UNet(nn.Module):
    def __init__(self):
        super(UNet, self).__init__()
        # U-Net的左半边部分:卷积-激活-池化
        self.layer1 = nn.Sequential(
            nn.Conv2d(1, 64, 3),
            nn.ReLU(),

            nn.Conv2d(64, 64, 3),
            nn.ReLU()
        )
        self.maxpool1 = nn.MaxPool2d(2, 2)

        self.layer2 = ConvLayer(in_channel=64, times=2)
        self.maxpool2 = nn.MaxPool2d(2, 2)

        self.layer3 = ConvLayer(in_channel=128, times=2)
        self.maxpool3 = nn.MaxPool2d(2, 2)

        self.layer4 = ConvLayer(in_channel=256, times=2)
        self.maxpool4 = nn.MaxPool2d(2, 2)

        self.layer5 = ConvLayer(in_channel=512, times=2)

        # U-Net的右半边部分:上采样, 裁剪and拼接, 卷积-激活
        self.upconv1 = nn.ConvTranspose2d(1024, 512, 2, stride=2)
        self.uplayer1 = ConvLayer(in_channel=1024, times=0.5)

        self.upconv2 = nn.ConvTranspose2d(512, 256, 2, stride=2)
        self.uplayer2 = ConvLayer(in_channel=512, times=0.5)

        self.upconv3 = nn.ConvTranspose2d(256, 128, 2, stride=2)
        self.uplayer3 = ConvLayer(in_channel=256, times=0.5)

        self.upconv4 = nn.ConvTranspose2d(128, 64, 2, stride=2)
        self.uplayer4 = ConvLayer(in_channel=128, times=0.5)

        # 1*1卷积
        self.conv1_1 = nn.Conv2d(64, 2, 1)

    def forward(self, x):
        # U-Net左半边部分
        x1 = self.layer1(x)
        x1_1 = self.maxpool1(x1)

        x2 = self.layer2(x1_1)
        x2_1 = self.maxpool2(x2)

        x3 = self.layer3(x2_1)
        x3_1 = self.maxpool3(x3)

        x4 = self.layer4(x3_1)
        x4_1 = self.maxpool4(x4)

        x5 = self.layer5(x4_1)

        # U-Net右半边部分
        y1 = self.upconv1(x5)
        x4_cat = x4[:, :, 4:-4, 4:-4]  # copy and crop
        y1_cat = torch.cat((y1, x4_cat), dim=1)  # 拼接
        y1_1 = self.uplayer1(y1_cat)

        y2 = self.upconv2(y1_1)
        x3_cat = x3[:, :, 16:-16, 16:-16]  # copy and crop
        y2_cat = torch.cat((y2, x3_cat), dim=1)
        y2_1 = self.uplayer2(y2_cat)

        y3 = self.upconv3(y2_1)
        x2_cat = x2[:, :, 40:-40, 40:-40]  # copy and crop
        y3_cat = torch.cat((y3, x2_cat), dim=1)
        y3_1 = self.uplayer3(y3_cat)

        y4 = self.upconv4(y3_1)
        x1_cat = x1[:, :, 88:-88, 88:-88]  # copy and crop
        y4_cat = torch.cat((y4, x1_cat), dim=1)
        y4_1 = self.uplayer4(y4_cat)

        # 1*1卷积
        y5 = self.conv1_1(y4_1)
        return y5


class ConvLayer(nn.Module):  # 卷积层, 封装简化计算, times为维度放缩倍数
    def __init__(self, in_channel, times, kernel_size=3):
        super(ConvLayer, self).__init__()
        self.l1 = nn.Sequential(
            nn.Conv2d(in_channel, int(times * in_channel), kernel_size),
            nn.ReLU()
        )

        self.l2 = nn.Sequential(
            nn.Conv2d(int(times * in_channel), int(times * in_channel), kernel_size),
            nn.ReLU()
        )

    def forward(self, x):
        x = self.l1(x)
        x = self.l2(x)
        return x


if __name__ == '__main__':
    inputs = torch.randn(10, 1, 572, 572)
    model = UNet()
    print(model)
    outputs = model(inputs)
    print("输入维度:", inputs.shape)
    print("输出维度:", outputs.shape)

我的疑问：
我没有用深度学习做过语义分析，以及图像分割任务，所以有的细节不是很明白，如果有了解的同学能够答疑，非常感谢。

1. U-Net的标签，是像素级的吗？每个像素点都有对应的标签，大小和整幅图像相同
2. U-Net的输出为2×388×388，输入是单通道572×572的灰度图
3. 第一个问题就是，输出图缩小了，和输入大小不同，那么像素级的标签还能一一对应吗
4. 第二个问题就是，一个标签，输出1×388×388还能够理解，与标签计算损失然后反向传播更新误差，2×388×388是代表有2个label矩阵吗，分别代表什么

我的猜测：

1. 输出有2个特征图，就对应着2个标签矩阵，比如二值分割的结果（背景白色，物体黑色）以及语义分割（像素级）的标签
2. 大小为388×388，后面可以通过变换到等同输入图像的大小

欢迎交流

---

问题解答---------------------------分割线 ---------------------------时间：2021年11月07日10:24:37

参考1：作者-LoveMIss-Y【个人整理】语义分割网络U-Net的设计架构与设计思想
参考2：作者-Kanny广小隶-用于语义分割的U-Net为什么这么强？

参考博客1和博客2以及其它资料及室友@Owenhhhh的帮助，现在对我上面提出的疑惑进行解答

疑问1解答：标签是像素级的，每个像素点都有一个标记类别，与原图大小相同

疑问3解答：
先看原文插图

翻译：重叠平铺策略，用于无缝分割任意大图像（此处为EM堆栈中神经元结构的分割）。黄色区域的分割预测需要蓝色区域内的图像数据作为输入。缺少的输入数据通过镜像进行推断

①策略1，上文参考博客1中写的非常清楚，我们可以知道，原图与标签可能都是388×388的，但是把原图经过上面的镜像翻转策略将图像变大一圈，达到572×572作为输入，然后在最后计算损失的时候只取扩充后图像的中间388×388部分，这样就大小一致啦

②策略2，把输出388×388得到的特征图，进行上采样与原图572×572保持一致

③策略3，设计网络的时候，进行padding填充，使得输入输出为等大小

疑问4解答：最后输出的特征图是有2个channel，是因为原图需要分割的图像包含两个类别（前景和后景）、需要n个类别就设计n+1个channel（n个类+1个背景），最后网络得到N个channel的特征图，将N个channel特征图对应相同位置的N个像素点做Softmax，得到N张概率矩阵（相同位置N个channel特征图像素点对应概率和为1），然后合并成一个预测类别图，最后合并的这张预测类别图，取当前位置最大概率所对应的类别，最后与标签进行损失计算。

上文中我的猜测不太对，以上就是我对U-Net的理解，对U-Net的学习就暂时告一段落，存在的问题还请各位指出！

今天是S11夺冠的11月7日，最后说一句EDG，牛B！

---

更新

2022年02月19日15:06:01语义分割任务，后面的损失是如何做的呢？
答：

• ①输入维度[3, 388, 388]；输出维度[n_class, 388, 388]，其中n_class为分类数（包含背景）；标签label大小为[1, 388, 388]同原图大小，需用标注工具对每一个像素打上标签。
• ②网络下采样进行特征提取，然后上采样得到输出[n_class, 388, 388]，输出值与标签值做损失，比如分成3类，输出维度为[3, 388, 388]，标签为[1, 388, 388]；简单理解就是将每一个像素点进行分类，一个像素点有3种可能性，对每一个像素点做softmax得到概率，然后用交叉熵计算损失。代码实现的时候，可以直接用交叉熵损失函数计算损失如下：

# 输出维度[N， 3， 388， 388],N为batch的维度
outputs = torch.rand(10, 3, 388, 388）
label = torch.rand(10, 1, 388, 388)
# 可以直接利用交叉熵损失函数计算损失
loss_softmax = nn.CrossEntropyLoss()
loss = (outputs, label)

过段时间会用自己实现的网络，利用U-Net进行语义分割。因为现在要写论文，所以需要等待一段时间。