unet图像分割pytorch代码

UNet是一种流行的卷积神经网络结构，专用于图像分割任务。在PyTorch中实现UNet需要导入必要的库，构建编码器解码器结构，并训练模型。下面是一个简化的UNet代码示例：

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
from torchvision.transforms import ToTensor

# 定义UNet的基本模块
class UNetBlock(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, dropout_rate=0.5):
        super(UNetBlock, self).__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(),
            nn.Dropout(dropout_rate)
        )
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(out_channels, in_channels // 2, kernel_size, stride, padding),
            nn.BatchNorm2d(in_channels // 2),
            nn.ReLU(),
            nn.Dropout(dropout_rate)
        )

    def forward(self, x):
        x1 = self.encoder(x)
        return self.decoder(x1)

# 全局定义UNet
class UNet(nn.Module):
    def __init__(self, num_classes):
        super(UNet, self).__init__()
        self.down1 = UNetBlock(3, 64) # 输入通道数为3 (RGB), 输出通道数为64
        self.pool1 = nn.MaxPool2d(2, 2)
        self.down2 = UNetBlock(64, 128)
        self.pool2 = nn.MaxPool2d(2, 2)
        self.down3 = UNetBlock(128, 256)
        self.pool3 = nn.MaxPool2d(2, 2)
        self.center = UNetBlock(256, 512)
        self.up3 = UNetBlock(512 + 256, 256)
        self.up2 = UNetBlock(256 + 128, 128)
        self.up1 = UNetBlock(128 + 64, 64)
        self.outconv = nn.Conv2d(64, num_classes, 1)

    def forward(self, x):
        x1 = self.down1(x)
        x2 = self.pool1(x1)
        x3 = self.down2(x2)
        x4 = self.pool2(x3)
        x5 = self.down3(x4)
        x = self.pool3(x5)
        x = self.center(x)
        
        x = torch.cat((x5, x), dim=1)
        x = self.up3(x)
        x = torch.cat((x4, x), dim=1)
        x = self.up2(x)
        x = torch.cat((x3, x), dim=1)
        x = self.up1(x)
        x = self.outconv(x)
        return x

# 示例用法
num_classes = 2 # 二分类问题，比如前景和背景
model = UNet(num_classes)
# 加载数据、设置优化器、损失函数等并开始训练...