U-Net是继FCN之后又一个经典的语义分割网络模型，并且也是很多后续语义分割模型的“祖宗”。这个网络模型是2015年提出来的，它具有一个非常对称的结构，很像字母“U”，所以被称作U-Net。U-Net被广泛应用于医学影像领域，如分割器官、肿瘤、血管等。由于医学图像通常具有复杂的结构和较高的分辨率要求，U-Net 的结构特点使其能够有效地处理这些图像，并提供准确的分割结果，为医学诊断和治疗提供重要的支持，在遥感卫星分析和工业质量检测中也有广泛的应用。U-Net相比于FCN，它对于细节的把控能力更强。U-Net总体结构如下：

和FCN一样，它也分成了编码和解码模块两部分，编码模块也被称作下采样，解码模块也被称作上采样。
1.输入572x572，经过3x3卷积和ReLU，大小变为570x570。（卷积计算公示：output = (input-kernel_size+2*padding)/stride + 1），再做一次卷积和ReLU，大小变为568x568，通道数变成64；
2.经过最大池化层，长宽变为原来的一半，变成284x284，经过两次卷积和ReLU，大小变成280x280，通道数变成128；
3.再次经过最大池化层，长宽变为原来的一半，变成140下140，再经过两次卷积和ReLU，大小变成136x136，通道数加倍变成256；
4.再次经过最大池化层，长宽继续减半，变成68x68，再经过两次卷积和ReLU，大小变成64x64，通道数变成512；
5.再次经过最大池化层，长宽继续减半，变成32x32，再经过两次卷积和ReLU，大小变成28x28，通道数变成1024；
6.此处开始上采样。经过2x2上采样，长宽变成原来的两倍，56x56，同时通道数减半，变成512;
7.从下采样过程中的512x64x64结果中，裁剪512x56x56的部分，和上采样后得到的512x56x56结果，进行通道维度上的拼接，注意不是求和是拼接concat，得到1024x56x56的结果，再做一次卷积和ReLU，通道数减半，变成512x54x54，再做一次卷积和ReLU，这次通道数保持不变，大小变成了52x52，所以这一层最终输出512x52x52；
8.上采样，通道数进一步减小变成256x256，长宽再变大一倍，得到256x104x104，和下采样过程中通道是256x136x136的部分拼接，因为长宽不同，所以需要裁剪，将下采样过程中的结果裁剪成256x104x104，按通道维度拼接，得到512x104x104，再经过卷积和ReLU，通道数减半，变成256x102x102，再经过卷积和ReLU，得到256x100x100；
9.上采样，通道数进一步减半变成128x128，长宽再变大一倍，得到128x200x200，和下采样过程中通道是128的部分拼接，裁剪成200x200大小后，按通道维度拼接，得到256x200x200，经过卷积和ReLU，通道数减半，得到128x198x198，再经过一次卷积和ReLU，通道不变，得到128x196x196；
10.上采样，通道数进一步减半变成64x64，长宽再变大一倍，得到64x392x392，和下采样过程中通道是64的部分拼接，裁剪成392x392大小后，按通道维度拼接，得到128x392x392，经过卷积和ReLU，通道数减半，得到64x390x390，再经过一次卷积和ReLU，通道不变，得到64x388x388，最后经过一次1x1卷积，改变通道数为2，得到2x388x388，也就是最终的输出。论文最终输出是分成了两类。
不过这里有一个问题，就是输入图像的大小和输出的大小是不一致的。输入是572X572的图像，输出的结果是388X388的，当然在医学影像中可能是无所谓的，不过在大部分的语义分割场景，最好还是使得输入图像和输出结果的大小是一致的，大部分的语义分割数据集，如VOC2012的图像和标签大小也是一直的，所以我们最好使得输入和输出大小一致。
U-Net的简单实现如下：

import torch.nn as nn
import torchvision.models as models
import torch

class UNetEncoder(nn.Module):
    def __init__(self, pretrained=True) -> None:
        super(UNetEncoder, self).__init__()
        vgg = models.vgg16(pretrained=pretrained)
        self.encoder = nn.Sequential(*list(vgg.children())[:-2])
        self.layer1 = nn.Sequential(vgg.features[:4])
        self.layer2 = nn.Sequential(vgg.features[4:9])
        self.layer3 = nn.Sequential(vgg.features[9:16])
        self.layer4 = nn.Sequential(vgg.features[16:23])
        self.layer5 = nn.Sequential(vgg.features[23:-1])

    def forward(self, x):
        features = []
        x = self.layer1(x)
        print("x.shape: ",x.shape)
        features.append(x)
        x = self.layer2(x)
        print("x.shape: ",x.shape)
        features.append(x)
        x = self.layer3(x)
        print("x.shape: ",x.shape)
        features.append(x)
        x = self.layer4(x)
        print("x.shape: ",x.shape)
        features.append(x)
        x = self.layer5(x)
        print("x.shape: ",x.shape)
        features.append(x)
        return features
    
class UNetDecoder(nn.Module):
    def __init__(self, num_classes) -> None:
        super(UNetDecoder, self).__init__()
        # 定义解码器层
        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        self.conv1 = nn.Conv2d(in_channels=1024, out_channels=512, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(in_channels=768, out_channels=256, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(in_channels=384, out_channels=128, kernel_size=3, padding=1)
        self.conv4 = nn.Conv2d(in_channels=192, out_channels=64, kernel_size=3, padding=1)
        self.classifier = nn.Conv2d(in_channels=64, out_channels=num_classes, kernel_size=1)

    def forward(self, features):
        x5 = features[-1] # 最后一层编码的结果
        print("x5.shape: ",x5.shape)
        x4 = self.upsample(x5)
        x4 = torch.cat([x4, features[-2]], dim=1) # 根据通道维度进行合并，第0维是batch_size
        x4 = self.conv1(x4)
        print("x4.shape: ",x4.shape)
        x3 = self.upsample(x4)
        x3 = torch.cat([x3, features[-3]], dim=1)
        x3 = self.conv2(x3)
        print("x3.shape: ",x3.shape)
        x2 = self.upsample(x3)
        x2 = torch.cat([x2, features[-4]], dim=1)
        x2 = self.conv3(x2)
        print("x2.shape: ",x2.shape)
        x1 = self.upsample(x2)
        x1 = torch.cat([x1, features[-5]], dim=1)
        x1 = self.conv4(x1)
        print("x1.shape: ",x1.shape)
        out = self.classifier(x1)
        return out
    
class UNet(nn.Module):
    def __init__(self, num_classes, pretrained=True):
        super(UNet, self).__init__()
        self.encoder = UNetEncoder(pretrained=pretrained)
        self.decoder = UNetDecoder(num_classes)

    def forward(self, x):
        features = self.encoder(x)
        out = self.decoder(features)
        return out
    
if __name__ == '__main__':
    model = UNet(num_classes=21, pretrained=True)
    x = torch.randn(1, 3, 480, 320)
    out = model(x)
    print(out.shape)
# 输出：
x.shape:  torch.Size([1, 64, 480, 320])
x.shape:  torch.Size([1, 128, 240, 160])
x.shape:  torch.Size([1, 256, 120, 80])
x.shape:  torch.Size([1, 512, 60, 40])
x.shape:  torch.Size([1, 512, 30, 20]) 
x5.shape:  torch.Size([1, 512, 30, 20])
x4.shape:  torch.Size([1, 512, 60, 40])
x3.shape:  torch.Size([1, 256, 120, 80])
x2.shape:  torch.Size([1, 128, 240, 160])
x1.shape:  torch.Size([1, 64, 480, 320])
torch.Size([1, 21, 480, 320])

可以看到，输入是四维的数据，batch_size=1，三通道的彩色图像，宽480，高320的一幅图像，输出batch_size=1，通道数21，用来分类，宽480，高320的分割结果图像。
整个U-Net代码分为编码器Encoder和解码器Decoder。Encoder就是下采样模块，采用预训练的vgg模型的特征提取模块。Decoder就是上采样模块，把特征图不断恢复为原始图像大小，并在整个过程中，按通道维度拼接特征提取过程中的特征图。
根据输出可以看到图像的变换过程，如下图所示：

下面，我们来看看这个U-Net模型在GID和VOC2012两个数据集上的分割效果：