UNet 图像分割

UNet 是一种专门用于图像分割的深度学习网络。它的核心思想是 U 型编码-解码结构 + 跳跃连接（Skip Connection）。 不同于图像分类（给整张图一个标签），图像分割需要给每个像素分配一个类别标签。

类比：想象你在看一张照片。先把它缩小成缩略图了解全局（编码器），然后再逐像素放大精确定位每个物体的边界（解码器）。 跳跃连接就像在放大的过程中，随时回头看原始照片的细节，确保不丢失边缘信息。

诞生背景

UNet 由 Olaf Ronneberger 等人于 2015 年在论文 "U-Net: Convolutional Networks for Biomedical Image Segmentation" 中提出， 最初是为医学图像分割设计的。医学图像的特点是训练数据少、需要精确的像素级分割，UNet 的跳跃连接正好解决了下采样过程中细节丢失的问题。

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


# 基础卷积块：两次 3×3 卷积 + BN + ReLU
class DoubleConv(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
        )

    def forward(self, x):
        return self.double_conv(x)


# 下采样块：MaxPool + DoubleConv
class Down(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.maxpool_conv = nn.Sequential(
            nn.MaxPool2d(2),
            DoubleConv(in_ch, out_ch),
        )

    def forward(self, x):
        return self.maxpool_conv(x)


# 上采样块：转置卷积 + 跳跃连接拼接 + DoubleConv
class Up(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        # 转置卷积：通道减半，尺寸翻倍
        self.up = nn.ConvTranspose2d(in_ch, in_ch // 2, kernel_size=2, stride=2)
        self.conv = DoubleConv(in_ch, out_ch)  # 拼接后通道数 = in_ch

    def forward(self, x, skip):
        # 上采样
        x = self.up(x)
        # 跳跃连接：拼接编码器特征
        x = torch.cat([skip, x], dim=1)
        return self.conv(x)