U-Net 模型教程(代码:可直接运行)

U-Net 是一种用于图像分割的深度学习架构，广泛应用于医学图像分割等领域。其关键特点是使用了对称的编码器-解码器结构，通过跳跃连接（skip connections）将编码器部分的特征图传递给解码器部分，帮助恢复图像的细节。本文将带你一起构建一个简单的 U-Net 模型，并逐步解释其组成部分和代码实现。

1. U-Net 模型的基本结构

U-Net 模型由以下几部分组成：

• Contracting Path（编码器）：由多个卷积块（ConvBlock）组成，负责提取图像的特征。每个卷积块通常包含两个卷积层，之后使用池化层（MaxPooling）来降低空间分辨率。

• Expanding Path（解码器）：由多个上采样块（UpsamplingBlock）组成，通过反卷积（ConvTranspose2d）将特征图上采样到更大的尺寸，并将编码器的特征图（通过跳跃连接）融合到解码器中，以便恢复空间分辨率和细节。

• 最终卷积层：在网络的末尾使用卷积层将特征图映射到所需的输出类别数目。

2. U-Net 代码实现

2.1 U-Net 模型

首先，我们定义了 U-Net 模型的主体结构，包括编码器、解码器以及最终的卷积层。

import torch
import torch.nn as nn

class UNetModel(nn.Module):
    def __init__(self, input_channels=3, n_filters=32, n_classes=23):
        super(UNetModel, self).__init__()
        
        # Contracting Path (encoding)
        self.cblock1 = ConvBlock(input_channels, n_filters)
        self.cblock2 = ConvBlock(n_filters, n_filters * 2)
        self.cblock3 = ConvBlock(n_filters * 2, n_filters * 4)
        self.cblock4 = ConvBlock(n_filters * 4, n_filters * 8, dropout=0.3)
        self.cblock5 = ConvBlock(n_filters * 8, n_filters * 16, dropout=0.3, max_pooling=False)
        
        # Expanding Path (decoding)
        self.ublock6 = UpsamplingBlock(n_filters * 16, n_filters * 8, n_filters * 8)
        self.ublock7