当前位置：首页 > news >正文

【深度学习|PyTorch】基于 PyTorch 搭建 U-Net 深度学习语义分割模型——附代码及其解释！

news 2025/8/5 10:47:56

【深度学习|PyTorch】基于 PyTorch 搭建 U-Net 深度学习语义分割模型——附代码及其解释！

论文地址： https://arxiv.org/abs/1505.04597
代码地址：https://github.com/jakeret/tf_unet

文章目录

【深度学习|PyTorch】基于 PyTorch 搭建 U-Net 深度学习语义分割模型——附代码及其解释！
1.数据准备
2.模型搭建：U-Net
3.模型训练
4.模型评估
总结

1.数据准备

语义分割任务的输入通常是图像以及对应的像素级标签（即每个像素的分类）。我们首先需要加载和预处理数据。

import torch
from torch.utils.data import DataLoader, Dataset
import torchvision.transforms as transforms
from PIL import Image
import osclass SegmentationDataset(Dataset):def __init__(self, image_dir, mask_dir, transform=None):self.image_dir = image_dirself.mask_dir = mask_dirself.transform = transformself.images = os.listdir(image_dir)def __len__(self):return len(self.images)def __getitem__(self, index):img_path = os.path.join(self.image_dir, self.images[index])mask_path = os.path.join(self.mask_dir, self.images[index])image = Image.open(img_path).convert("RGB")mask = Image.open(mask_path).convert("L")  # Assuming masks are grayscaleif self.transform is not None:image = self.transform(image)mask = self.transform(mask)return image, mask# 数据加载及预处理
image_dir = "path_to_images"
mask_dir = "path_to_masks"transform = transforms.Compose([transforms.Resize((256, 256)),transforms.ToTensor(),
])dataset = SegmentationDataset(image_dir, mask_dir, transform)
dataloader = DataLoader(dataset, batch_size=4, shuffle=True)

代码解释：

SegmentationDataset：自定义的数据集类，负责读取图像和对应的掩码文件（标签）。
__getitem__ 方法：从文件夹中加载图像和对应的掩码，并进行相应的预处理。
transforms：使用 torchvision 中的 transforms 对图像进行调整（例如，缩放和转换为 Tensor）。

2.模型搭建：U-Net

U-Net 是一种常用于医学图像分割的卷积神经网络。其结构包括下采样路径（编码器）和上采样路径（解码器），并在同一层级将特征图通过跳跃连接传递。

import torch.nn as nn
import torchclass UNet(nn.Module):def __init__(self, in_channels=3, out_channels=1):super(UNet, self).__init__()# Contracting path (Encoder)self.enc_conv1 = self.double_conv(in_channels, 64)self.enc_conv2 = self.double_conv(64, 128)self.enc_conv3 = self.double_conv(128, 256)self.enc_conv4 = self.double_conv(256, 512)# Maxpooling layerself.pool = nn.MaxPool2d(kernel_size=2, stride=2)# Expansive path (Decoder)self.up_conv3 = self.up_conv(512, 256)self.dec_conv3 = self.double_conv(512, 256)self.up_conv2 = self.up_conv(256, 128)self.dec_conv2 = self.double_conv(256, 128)self.up_conv1 = self.up_conv(128, 64)self.dec_conv1 = self.double_conv(128, 64)# Final output layerself.final_conv = nn.Conv2d(64, out_channels, kernel_size=1)def double_conv(self, in_channels, out_channels):return nn.Sequential(nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),nn.ReLU(inplace=True),nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),nn.ReLU(inplace=True))def up_conv(self, in_channels, out_channels):return nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2)def forward(self, x):# Encoderenc1 = self.enc_conv1(x)enc2 = self.enc_conv2(self.pool(enc1))enc3 = self.enc_conv3(self.pool(enc2))enc4 = self.enc_conv4(self.pool(enc3))# Decoderdec3 = self.up_conv3(enc4)dec3 = torch.cat((dec3, enc3), dim=1)dec3 = self.dec_conv3(dec3)dec2 = self.up_conv2(dec3)dec2 = torch.cat((dec2, enc2), dim=1)dec2 = self.dec_conv2(dec2)dec1 = self.up_conv1(dec2)dec1 = torch.cat((dec1, enc1), dim=1)dec1 = self.dec_conv1(dec1)# Outputreturn self.final_conv(dec1)# 实例化模型
model = UNet(in_channels=3, out_channels=1).to('cuda' if torch.cuda.is_available() else 'cpu')

代码解释：

double_conv：U-Net 结构中每层包含两个卷积层，卷积核大小为3，使用 ReLU 激活函数。
up_conv：用于上采样的转置卷积层。
forward：定义了模型的前向传播路径，使用了 U-Net 的跳跃连接，保证上采样时能够使用对应层级的特征图。

3.模型训练

训练模型需要定义损失函数和优化器。我们通常使用交叉熵损失或者 Dice 损失进行语义分割任务。

import torch.optim as optim
import torch.nn.functional as F# 损失函数和优化器
criterion = nn.BCEWithLogitsLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-4)# 训练循环
num_epochs = 20
device = 'cuda' if torch.cuda.is_available() else 'cpu'for epoch in range(num_epochs):model.train()running_loss = 0.0for images, masks in dataloader:images = images.to(device)masks = masks.to(device)# Forward passoutputs = model(images)loss = criterion(outputs, masks)# Backward pass and optimizationoptimizer.zero_grad()loss.backward()optimizer.step()running_loss += loss.item()print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(dataloader)}")

代码解释：

criterion：使用二元交叉熵损失（BCEWithLogitsLoss）处理二分类分割任务。对于多类分割，可使用 CrossEntropyLoss。
optimizer：Adam 优化器，学习率设为 1e-4。
训练循环：每个 epoch 中，模型进行前向传播、计算损失、反向传播并更新权重。

4.模型评估

为了评估模型性能，可以使用常见的分割指标如 IoU（交并比）或 Dice 系数。

def dice_coefficient(preds, labels, threshold=0.5):preds = torch.sigmoid(preds)  # Apply sigmoid to get probabilitiespreds = (preds > threshold).float()  # Threshold predictionsintersection = (preds * labels).sum()union = preds.sum() + labels.sum()dice = 2 * intersection / (union + 1e-8)  # Add small epsilon to avoid division by zeroreturn dice# 在训练完成后，评估模型
model.eval()
with torch.no_grad():dice_score = 0.0for images, masks in dataloader:images = images.to(device)masks = masks.to(device)outputs = model(images)dice_score += dice_coefficient(outputs, masks)dice_score /= len(dataloader)print(f"Dice Coefficient: {dice_score}")