浅谈博星卓越网站建设,网站制作工具推荐,朔州网站建设费用,网站优化排名易下拉软件【深度学习|PyTorch】基于 PyTorch 搭建 U-Net 深度学习语义分割模型——附代码及其解释#xff01;
【深度学习|PyTorch】基于 PyTorch 搭建 U-Net 深度学习语义分割模型——附代码及其解释#xff01;
论文地址#xff1a; https://arxiv.org/abs/1505.04597 代码地址
【深度学习|PyTorch】基于 PyTorch 搭建 U-Net 深度学习语义分割模型——附代码及其解释
论文地址 https://arxiv.org/abs/1505.04597 代码地址https://github.com/jakeret/tf_unet 文章目录 【深度学习|PyTorch】基于 PyTorch 搭建 U-Net 深度学习语义分割模型——附代码及其解释1.数据准备2.模型搭建U-Net3.模型训练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, transformNone):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_maskstransform transforms.Compose([transforms.Resize((256, 256)),transforms.ToTensor(),
])dataset SegmentationDataset(image_dir, mask_dir, transform)
dataloader DataLoader(dataset, batch_size4, shuffleTrue)代码解释
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_channels3, out_channels1):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_size2, stride2)# 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_size1)def double_conv(self, in_channels, out_channels):return nn.Sequential(nn.Conv2d(in_channels, out_channels, kernel_size3, padding1),nn.ReLU(inplaceTrue),nn.Conv2d(out_channels, out_channels, kernel_size3, padding1),nn.ReLU(inplaceTrue))def up_conv(self, in_channels, out_channels):return nn.ConvTranspose2d(in_channels, out_channels, kernel_size2, stride2)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), dim1)dec3 self.dec_conv3(dec3)dec2 self.up_conv2(dec3)dec2 torch.cat((dec2, enc2), dim1)dec2 self.dec_conv2(dec2)dec1 self.up_conv1(dec2)dec1 torch.cat((dec1, enc1), dim1)dec1 self.dec_conv1(dec1)# Outputreturn self.final_conv(dec1)# 实例化模型
model UNet(in_channels3, out_channels1).to(cuda if torch.cuda.is_available() else cpu)代码解释
double_convU-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(), lr1e-4)# 训练循环
num_epochs 20
device cuda if torch.cuda.is_available() else cpufor 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(fEpoch [{epoch1}/{num_epochs}], Loss: {running_loss/len(dataloader)})代码解释
criterion使用二元交叉熵损失BCEWithLogitsLoss处理二分类分割任务。对于多类分割可使用 CrossEntropyLoss。optimizerAdam 优化器学习率设为 1e-4。训练循环每个 epoch 中模型进行前向传播、计算损失、反向传播并更新权重。
4.模型评估
为了评估模型性能可以使用常见的分割指标如 IoU交并比或 Dice 系数。
def dice_coefficient(preds, labels, threshold0.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(fDice Coefficient: {dice_score})代码解释
dice_coefficient计算 Dice 系数衡量预测和真实标签的重合程度值越接近 1 表示预测效果越好。评估模型时使用 model.eval() 关闭 dropout 等不影响推理过程的操作并使用 torch.no_grad() 以节省内存。
总结
以上是从数据准备、模型搭建、训练到精度评估的完整流程。我们基于 PyTorch 实现了一个 U-Net 语义分割模型并详解了每步的代码。
