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两个数据集上的分割效果：
从GID的分割结果可以看到U-Net的分割结果更为精细，细节部分比FCN提高了不少。
在VOC2012数据集上分割效果还比较一般，只能分割出大致形态，细节还不够完善。
这里有一个要注意的问题，就是U-Net并不是任意大小的输入都可以运行的，对于一些长宽不符合要求的会报错。因为在下采样过程中要进行卷积运算，由于长宽不一定是偶数，可能造成图像长宽变化，上采样后无法和下采样过程中的特征图进行拼接，所以训练过程中最好对图像进行裁剪操作，裁剪成适合U-Net运行的大小，或者改进U-Net模型，使其可以兼容各种不同大小的数据。