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

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

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

论文地址: https://arxiv.org/abs/1505.04597

代码地址:https://github.com/jakeret/tf_unet


文章目录


1.数据准备

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

clike 复制代码
import torch
from torch.utils.data import DataLoader, Dataset
import torchvision.transforms as transforms
from PIL import Image
import os

class SegmentationDataset(Dataset):
    def __init__(self, image_dir, mask_dir, transform=None):
        self.image_dir = image_dir
        self.mask_dir = mask_dir
        self.transform = transform
        self.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 grayscale

        if 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 是一种常用于医学图像分割的卷积神经网络。其结构包括下采样路径(编码器)和上采样路径(解码器),并在同一层级将特征图通过跳跃连接传递。

clike 复制代码
import torch.nn as nn
import torch

class 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 layer
        self.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 layer
        self.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):
        # Encoder
        enc1 = 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))

        # Decoder
        dec3 = 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)

        # Output
        return 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 损失进行语义分割任务。

clike 复制代码
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.0
    for images, masks in dataloader:
        images = images.to(device)
        masks = masks.to(device)

        # Forward pass
        outputs = model(images)
        loss = criterion(outputs, masks)

        # Backward pass and optimization
        optimizer.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 系数。

clike 复制代码
def dice_coefficient(preds, labels, threshold=0.5):
    preds = torch.sigmoid(preds)  # Apply sigmoid to get probabilities
    preds = (preds > threshold).float()  # Threshold predictions
    intersection = (preds * labels).sum()
    union = preds.sum() + labels.sum()
    dice = 2 * intersection / (union + 1e-8)  # Add small epsilon to avoid division by zero
    return dice

# 在训练完成后,评估模型
model.eval()
with torch.no_grad():
    dice_score = 0.0
    for 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}")

代码解释:

  • dice_coefficient:计算 Dice 系数,衡量预测和真实标签的重合程度,值越接近 1 表示预测效果越好。
  • 评估模型时使用 model.eval() 关闭 dropout 等不影响推理过程的操作,并使用 torch.no_grad() 以节省内存。

总结

以上是从数据准备、模型搭建、训练到精度评估的完整流程。我们基于 PyTorch 实现了一个 U-Net 语义分割模型,并详解了每步的代码。

相关推荐
Elastic 中国社区官方博客18 分钟前
使用 Elastic AI Assistant for Search 和 Azure OpenAI 实现从 0 到 60 的转变
大数据·人工智能·elasticsearch·microsoft·搜索引擎·ai·azure
江_小_白1 小时前
自动驾驶之激光雷达
人工智能·机器学习·自动驾驶
yusaisai大鱼3 小时前
TensorFlow如何调用GPU?
人工智能·tensorflow
麻瓜也要学魔法3 小时前
链路状态路由协议-OSPF
网络
Estar.Lee3 小时前
查手机号归属地免费API接口教程
android·网络·后端·网络协议·tcp/ip·oneapi
傻啦嘿哟4 小时前
代理IP在后端开发中的应用与后端工程师的角色
网络·网络协议·tcp/ip
湫ccc4 小时前
《Python基础》之字符串格式化输出
开发语言·python
Red Red4 小时前
网安基础知识|IDS入侵检测系统|IPS入侵防御系统|堡垒机|VPN|EDR|CC防御|云安全-VDC/VPC|安全服务
网络·笔记·学习·安全·web安全
mqiqe4 小时前
Python MySQL通过Binlog 获取变更记录 恢复数据
开发语言·python·mysql