原始U-Net模型代码

论文：1505.U-Net: Convolutional Networks for Biomedical Image Segmentation
代码： https://github.com/milesial/Pytorch-UNet/blob/master/unet/unet_model.py

一、原始u-net 架构结构 （输入572x572x1，5层，向下采样4次）：

每个蓝色框对应一个多通道特征图。通道的数量在框的顶部表示。x-y 大小在框的左下角提供。白盒代表复制的特征图。箭头表示不同的操作

1.0 框架代码

import torch
import torch.nn as nn

class UNet(nn.Module):
    def __init__(self, n_channels, n_classes, bilinear=False):
        """
        初始化函数，定义UNet模型的结构。

        参数:
        n_channels -- 输入图像的通道数，例如RGB图像的通道数为3。
        n_classes -- 输出的类别数，即分割后的图像通道数。
        bilinear -- 是否使用双线性插值进行上采样。默认值为False。
        """
        super(UNet, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.bilinear = bilinear

        # 初始卷积层，输入通道数为n_channels，输出通道数为64
        self.inc = DoubleConv(n_channels, 64)
        # 第一个下采样层，输入通道数为64，输出通道数为128
        self.down1 = Down(64, 128)
        # 第二个下采样层，输入通道数为128，输出通道数为256
        self.down2 = Down(128, 256)
        # 第三个下采样层，输入通道数为256，输出通道数为512
        self.down3 = Down(256, 512)
        # 根据是否使用双线性插值设置下采样因子
        factor = 2 if bilinear else 1
        # 第四个下采样层，输入通道数为512，输出通道数为1024（如果使用双线性插值，则输出通道数减半）
        self.down4 = Down(512, 1024 // factor)
        # 第一个上采样层，输入通道数为1024（或512），输出通道数为512（或256）
        self.up1 = Up(1024, 512 // factor, bilinear)
        # 第二个上采样层，输入通道数为512（或256），输出通道数为256（或128）
        self.up2 = Up(512, 256 // factor, bilinear)
        # 第三个上采样层，输入通道数为256（或128），输出通道数为128（或64）
        self.up3 = Up(256, 128 // factor, bilinear)
        # 第四个上采样层，输入通道数为128，输出通道数为64
        self.up4 = Up(128, 64, bilinear)
        # 输出卷积层，输入通道数为64，输出通道数为n_classes
        self.outc = OutConv(64, n_classes)

    def forward(self, x):
        """
        前向传播函数，定义输入x如何通过各层传递并输出结果。

        参数:
        x -- 输入的图像张量

        返回:
        logits -- 输出的类别概率张量
        """
        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)   # 第四个上采样层，并与对应的下采样层输出拼接
        logits = self.outc(x) # 输出卷积层
        return logits         # 返回最终的类别概率张量

1.1 DoubleConv：2个3x3卷积+relu层的实现

import torch
import torch.nn as nn
import torch.nn.functional as F

class DoubleConv(nn.Module):
    """(卷积 => [批归一化] => ReLU) * 2"""

    def __init__(self, in_channels, out_channels, mid_channels=None):
        """
        初始化函数，定义双卷积层结构

        参数:
        in_channels -- 输入通道数
        out_channels -- 输出通道数
        mid_channels -- 中间层通道数，如果未指定，则等于输出通道数
        """
        super().__init__()
        if not mid_channels:
            mid_channels = out_channels
        # 定义双卷积层，包含两次卷积，每次卷积后跟批归一化和ReLU激活函数
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(mid_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        """
        前向传播函数，输入x，输出双卷积的结果
        """
        return self.double_conv(x)

1.2 Down：向下采样+同层2次卷积 （特征图维度变为原来的一半）

每次向下采样后，经过2次卷积层

对应在架构图为

class Down(nn.Module):
    """通过最大池化下采样，然后进行双卷积"""

    def __init__(self, in_channels, out_channels):
        """
        初始化函数，定义下采样结构

        参数:
        in_channels -- 输入通道数
        out_channels -- 输出通道数
        """
        super().__init__()
        # 定义下采样层，包含一个2x2最大池化层，然后是一个双卷积层
        self.maxpool_conv = nn.Sequential(
            nn.MaxPool2d(2),
            DoubleConv(in_channels, out_channels)
        )

    def forward(self, x):
        """
        前向传播函数，输入x，输出下采样的结果
        """
        return self.maxpool_conv(x)

1.3 Up: 向上采样+拼接（特征图维度加倍+拼接之前维度）

每次向上采样后，要接收对应向下采样结果的特征图拼接（padding和crop是类似的方法）
拼接过程中，对称层效果最好

拼接完成后，经过2次卷积处理）

代码中up操作

class Up(nn.Module):
    """通过上采样，然后进行双卷积"""

    def __init__(self, in_channels, out_channels, bilinear=True):
        """
        初始化函数，定义上采样结构

        参数:
        in_channels -- 输入通道数
        out_channels -- 输出通道数
        bilinear -- 是否使用双线性插值进行上采样，默认值为True
        """
        super().__init__()

        # 如果使用双线性插值，则使用常规卷积来减少通道数
        if bilinear:
            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
            self.conv = DoubleConv(in_channels, out_channels