pytorch复现u-net

### 使用 PyTorch 实现 U-Net模型，以下是详细的代码示例[^1]。

#### 导入必要的库

import torch
import torch.nn as nn
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, Dataset

#### 定义双卷积模块
此部分定义了一个基础组件，该组件由两个连续的3×3卷积层组成，每层后面跟着批标准化和ReLU激活函数。

class DoubleConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(DoubleConv, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

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

#### 构建完整的 U-Net 结构
整个网络结构遵循编码器-解码器架构，在下采样路径中逐步减少空间维度的同时增加通道数；而在上采样路径则相反操作，并最终输出与输入尺寸相同的预测图。

class UNet(nn.Module):
    def __init__(self, n_class):
        super(UNet, self).__init__()

        # 下采样/收缩路径
        self.down_conv_1 = DoubleConv(3, 64)  # 输入RGB三通道图片
        self.down_conv_2 = DoubleConv(64, 128)
        self.down_conv_3 = DoubleConv(128, 256)

        # 中间桥接层
        self.double_conv = DoubleConv(256, 512)

        # 上采样/扩展路径
        self.up_transpose_1 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)
        self.up_conv_1 = DoubleConv(512, 256)

        self.up_transpose_2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
        self.up_conv_2 = DoubleConv(256, 128)

        self.out = nn.Conv2d(128, n_class, kernel_size=1)  # 输出类别数量

    def forward(self, image):
        # 编码过程
        x1 = self.down_conv_1(image)
        x2 = F.max_pool2d(x1, kernel_size=2, stride=2)
        x3 = self.down_conv_2(x2)
        x4 = F.max_pool2d(x3, kernel_size=2, stride=2)
        x5 = self.down_conv_3(x4)
        x6 = F.max_pool2d(x5, kernel_size=2, stride=2)

        # 底部瓶颈连接
        x7 = self.double_conv(x6)

        # 解码过程
        x = self.up_transpose_1(x7)
        y = crop_img(x5, x.shape[2:])  # 对应裁剪跳跃连接处feature map大小一致
        x = self.up_conv_1(torch.cat([x, y], dim=1))

        x = self.up_transpose_2(x)
        y = crop_img(x3, x.shape[2:])
        x = self.up_conv_2(torch.cat([x, y], dim=1))
        
        output = self.out(x)
        return output

请注意上述`crop_img()`函数未在此展示，其作用是在拼接过程中调整特征图尺寸匹配。实际应用时还需要考虑数据增强、损失函数的选择等问题[^3]。