FEATURED · 精选文章

PyTorch实现语义分割:UNet与DeepLabV3核心差异与代码解析

发布时间 / 2026/9/13 16:48:38
来源 / 创域科博编辑部
栏目 / 资讯中心
PyTorch实现语义分割:UNet与DeepLabV3核心差异与代码解析 简介基于Pytorch的语义分割算法实现合集集成DeepLabV3、UNet、FCN8s等主流网络适用于计算机视觉学习者与算法工程人员快速搭建分割模型、开展训练与测试。资源已提供训练好的DeepLabV3Xception和FCN8s权重文件同时包含测试输出图片可直接用于效果验证与结果对比。代码按backbones、datasets、models、utils等模块划分数据集csv管理、数据增强、损失函数、日志计时及评价指标均已封装目录结构清晰便于二次开发与学习。包内共157个文件以40个Python源码文件为核心辅以112张测试结果图片、2份项目说明Markdown文档和少量图片素材整体压缩包仅75.96MB轻量易下载。已有839人浏览/学习代码均经过运行测试适合在公开数据集上做语义分割入门实践或课题复现。1. 从一只猫身上看懂语义分割该学DeepLabV3还是UNet很多人一上来就下载一套开源项目跑通UNet就以为语义分割学会了。真到了实际场景——自动抠图、遥感地物分类、医疗影像肿瘤勾画——换个数据集就连loss曲线变平是收敛还是梯度消失都分不清。问题不在参数在于没吃透模型之间的结构差异。DeepLabV3和UNet恰好代表了语义分割的两条主流路线一个靠空洞卷积扩大感受野搞定多尺度目标一个靠编码器解码器和跳跃连接把细节从浅层一路带回来。用PyTorch把这两个模型从零写出来而不是直接调torchvision里的封装才能理解每个跳连、每个膨胀率的取舍。这篇就顺着数据集、网络实现、训练评估这条线把两套方案的代码和选型逻辑一次讲透。2. 用PyTorch搭建环境并制作语义分割数据集2.1 PyTorch版本与CUDA匹配装错了后续全白搭拿到Python源码先别急着跑环境不匹配是pytorch安装阶段最典型的坑。PyTorch 2.x推荐Python 3.8-3.11CUDA版本要和显卡驱动配套。Anaconda管理环境是最省心的方式下面是在Linux和Windows下通用的安装流程conda create -n seg python3.9 -y conda activate seg # 先确认显驱动支持的CUDA版本nvidia-smi查看右上角 pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118代码里两个参数值得说清楚cu118对应CUDA 11.8如果你的驱动是CUDA 12.x就换成cu121CPU训练的同学把--index-url去掉装默认CPU版本。版本选错最常见的现象是torch.cuda.is_available()返回False但这不代表安装失败多数是驱动太老或CUDA Runtime与PyTorch编译版本不一致。提示装完用python -c import torch;print(torch.__version__, torch.cuda.is_available())验证返回True再进入下一步。2.2 VOC/COCO格式与自定义数据集目录结构语义分割数据集怎么制作决定了训练脚本的适配成本。最通用的组织方式沿袭VOC格式图片放JPEGImages标签放SegmentationClass分割掩码是单通道的PNG每个像素值对应类别编号0一般是背景。data/ ├── JPEGImages/ # RGB原图 │ ├── img_001.jpg │ └── img_002.jpg ├── SegmentationClass/ # 单通道掩码像素值从0到N-1 │ ├── img_001.png │ └── img_002.png └── train.txt # 每一行写不带后缀的文件名自制数据集最常用的工具是Labelme画多边形后转成掩码。转换的核心逻辑是把JSON里的每个多边形用cv2.fillPoly()填充到一张全零的单通道图上。这里有一个非常容易踩的坑PIL打开掩码图时如果模式是P直接转成数组后像素值是调色板索引而不是原始类别编号必须np.array(img.convert(L))才行。COCO格式则走pycocotools它的标注是RLE游程编码训练时解码开销更大但类别多的时候标注成本低。二选一的话跑UNet和DeepLabV3的对比实验直接用VOC格式最省事后面封装Dataset的时候只需要管图片和标签两套路径映射。2.3 自定义Dataset类与数据增强的默认参数PyTorch的Dataset类和DataLoader配合决定了数据怎么进模型。核心是重写__getitem__在这个函数里完成图像和标签的同步读取、同步增强。下面是适合UNet和DeepLabV3共用的代码骨架import torch import numpy as np import cv2 from torch.utils.data import Dataset from torchvision import transforms class SegDataset(Dataset): def __init__(self, img_dir, mask_dir, file_list, crop_size(512, 512)): self.img_dir img_dir self.mask_dir mask_dir self.files file_list self.crop_size crop_size # 归一化参数ImageNet统计量迁移学习常用的均值标准差 self.img_transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) def __len__(self): return len(self.files) def __getitem__(self, idx): name self.files[idx].strip() img cv2.imread(f{self.img_dir}/{name}.jpg) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) mask cv2.imread(f{self.mask_dir}/{name}.png, cv2.IMREAD_GRAYSCALE) # 随机裁剪保证图片和标签裁剪到同一区域 h, w img.shape[:2] ch, cw min(self.crop_size[0], h), min(self.crop_size[1], w) top np.random.randint(0, h - ch 1) left np.random.randint(0, w - cw 1) img img[top:topch, left:leftcw] mask mask[top:topch, left:leftcw] img self.img_transform(img) mask_tensor torch.from_numpy(mask).long() # 标签必须是long型CrossEntropyLoss只接受这类 return img, mask_tensor这段代码里Normalize的均值和标准差是ImageNet预训练模型的标准配置如果从零训练也可以全填0。mask转torch.long()是新手最容易漏的一步CrossEntropyLoss内部做索引查找标签必须是整型。DataLoader配置一般设batch_size8显存8G以下用4、num_workers4、shuffleTrue这里的num_workers在Windows上偶尔会报BrokenPipe错误调成0能规避但训练速度会慢。3. UNet语义分割实现编码器、解码器与跳跃连接3.1 U型结构与跳跃连接要解决什么问题UNet最核心的洞察是语义信息在编码器里越深越抽象但分辨率越来越低直接上采样回去会丢失边缘细节。于是作者设计了跳跃连接把编码器每层下采样前的特征图直接拼到解码器对应分辨率的特征图后面。这样解码器既能看到深层的高层语义又能看到浅层的边缘纹理两个信息互补。这也是为什么UNet在医学图像上表现尤其好——分割目标的边界本身就模糊浅层细节权重高。在Pytorch里实现UNet最底层的积木是conv-bn-relu三元组。下采样用maxpool上采样用ConvTranspose2d。注意UNet的上采样不做跳跃连接而是在解码器内部逐级恢复分辨率这条信息流从最底层一路传到输出头。3.2 可运行的UNet核心代码import torch import torch.nn as nn import torch.nn.functional as F class DoubleConv(nn.Module): 每个U型块内部的两层卷积 def __init__(self, in_ch, out_ch): super().__init__() self.conv nn.Sequential( nn.Conv2d(in_ch, out_ch, kernel_size3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, kernel_size3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue) ) def forward(self, x): return self.conv(x) class UNet(nn.Module): def __init__(self, in_channels3, num_classes21, base_ch64): super().__init__() # 编码器四个阶段通道数倍增 self.enc1 DoubleConv(in_channels, base_ch) self.pool1 nn.MaxPool2d(2) self.enc2 DoubleConv(base_ch, base_ch * 2) self.pool2 nn.MaxPool2d(2) self.enc3 DoubleConv(base_ch * 2, base_ch * 4) self.pool3 nn.MaxPool2d(2) self.enc4 DoubleConv(base_ch * 4, base_ch * 8) self.pool4 nn.MaxPool2d(2) self.bottleneck DoubleConv(base_ch * 8, base_ch * 16) # 解码器四个阶段通道数递减 self.up4 nn.ConvTranspose2d(base_ch * 16, base_ch * 8, kernel_size2, stride2) self.dec4 DoubleConv(base_ch * 16, base_ch * 8) self.up3 nn.ConvTranspose2d(base_ch * 8, base_ch * 4, kernel_size2, stride2) self.dec3 DoubleConv(base_ch * 8, base_ch * 4) self.up2 nn.ConvTranspose2d(base_ch * 4, base_ch * 2, kernel_size2, stride2) self.dec2 DoubleConv(base_ch * 4, base_ch * 2) self.up1 nn.ConvTranspose2d(base_ch * 2, base_ch, kernel_size2, stride2) self.dec1 DoubleConv(base_ch * 2, base_ch) self.out nn.Conv2d(base_ch, num_classes, kernel_size1) def forward(self, x): # 编码器路径保存四个跳跃连接特征 e1 self.enc1(x) e2 self.enc2(self.pool1(e1)) e3 self.enc3(self.pool2(e2)) e4 self.enc4(self.pool3(e3)) bn self.bottleneck(self.pool4(e4)) # 解码器路径先上采样再拼接跳跃连接 d4 self.dec4(torch.cat([self.up4(bn), e4], dim1)) d3 self.dec3(torch.cat([self.up3(d4), e3], dim1)) d2 self.dec2(torch.cat([self.up2(d3), e2], dim1)) d1 self.dec1(torch.cat([self.up1(d2), e1], dim1)) return self.out(d1)torch.cat沿着dim1拼接要求两路特征的宽高严格一致否则直接报张量形状不匹配的错误。输入尺寸必须能被16整除因为网络里一共4次下采样512x512的输入会变成32x32的特征图。如果把base_ch从64改成32模型参数减少约四分之一显存占用下降但分割精度会轻微回落视觉任务里通道数不是越多越好。3.3 UNet改进的方向在哪里现在引擎里的UNet大多是原始结构的变体。最常见的unet模型改进是替换编码器为预训练的ResNet或EfficientNet这样能用上大规模数据集学到的特征。另一个常见做法是把ConvTranspose2d换成Upsample(modebilinear)再跟一个卷积能减少棋盘效应。FastSCNN和UNet则是在跳跃连接上做文章把简单拼接改成密集连接或者加注意力模块。跑通了上面这份代码之后可以先从残差块替换、BN换成GroupNorm两个方向做消融实验观察验证集mIoU的变化幅度比自己拍脑袋加模块可靠得多。4. DeepLabV3语义分割实现空洞卷积与ASPP4.1 空洞卷积为什么能在不降低分辨率的前提下扩大感受野DeepLab系列的核心在于空洞卷积。普通3x3卷积在stride为2的下采样层里损失了空间细节空洞卷积通过在卷积核内部插入0值不增加参数量的前提下把感受野撑大。膨胀率rate2时一个3x3卷积核覆盖的区域是5x5但有效参数仍然是9个。这样设计的好处是feature map尺寸不需要缩小在保留边缘信息的同时能聚合更大范围的上下文。DeepLabV3的网络结构图中主干是ResNet101最后两个stage将stride从16调整到8即去掉这部分的downsample层靠空洞卷积弥补减小的stride。这个微调是DeepLabV3复现成败的地方只把backbone换成ResNet不调整最后几个block的stride输出分辨率直接减半分割边缘变得一块一块的。4.2 DeepLabV3关键模块代码import torch import torch.nn as nn import torchvision.models as models class ASPP(nn.Module): 空洞空间金字塔池化多分支不同膨胀率 def __init__(self, in_ch2048, out_ch256): super().__init__() # 1x1卷积分支等价于rate1 self.conv1 nn.Sequential( nn.Conv2d(in_ch, out_ch, 1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue) ) # 三个不同膨胀率的3x3空洞卷积 self.conv2 nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding6, dilation6, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue) ) self.conv3 nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding12, dilation12, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue) ) self.conv4 nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding18, dilation18, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue) ) # 全局平均池化分支捕获全图上下文 self.pool nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Conv2d(in_ch, out_ch, 1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue) ) self.project nn.Sequential( nn.Conv2d(out_ch * 5, out_ch, 1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Dropout(0.1) ) def forward(self, x): x1 self.conv1(x) x2 self.conv2(x) x3 self.conv3(x) x4 self.conv4(x) x5 self.pool(x) # 池化分支要插值回原尺寸 x5 nn.functional.interpolate(x5, sizex4.shape[2:], modebilinear, align_cornersFalse) x torch.cat([x1, x2, x3, x4, x5], dim1) return self.project(x) class DeepLabV3(nn.Module): def __init__(self, num_classes21, backboneresnet50): super().__init__() # 使用torchvision提供的预训练模型修改fc层为空洞结构 self.backbone models.resnet50(weightsmodels.ResNet50_Weights.IMAGENET1K_V1) # 冻结前两层加速训练且减少过拟合 for p in self.backbone.layer1.parameters(): p.requires_grad False # 替换layer3和layer4的下采样换成空洞卷积 self.backbone.layer4.apply(self._replace_stride_with_dilation) self.aspp ASPP(in_ch2048, out_ch256) self.classifier nn.Conv2d(256, num_classes, kernel_size1) def _replace_stride_with_dilation(self, module): if isinstance(module, nn.Conv2d) and module.stride (2, 2): module.stride (1, 1) module.dilation (2, 2) module.padding (2, 2) def forward(self, x): x self.backbone.conv1(x) x self.backbone.bn1(x) x self.backbone.relu(x) x self.backbone.maxpool(x) x self.backbone.layer1(x) x self.backbone.layer2(x) x self.backbone.layer3(x) x self.backbone.layer4(x) x self.aspp(x) x self.classifier(x) return x_replace_stride_with_dilation这个函数是整个复现的关键。layer4的第一个卷积stride从2改成1会导致后续分辨率增大一倍对应地dilation从1改成2padding从1改成2两者平衡后输出张量的空间尺寸不变。如果漏了这步ASPP输入的维度对不上预训练权重的2048通道直接报错。冻结layer1参数能把显存占用压下去10%左右精度损失在1%以内。4.3 冻结策略、FrozenBN与多卡训练的小细节DeepLabV3迁移学习时torchvision的预训练Backbone默认会跑BatchNorm的batch级统计量小batch size下BN均值和方差估计不稳。一个常见做法是Backbone进入eval模式保持FrozenBN状态ASPP和解码头走train模式。实现上在训练循环里先model.backbone.eval()再model.aspp.train()、model.classifier.train()分开控制。多卡训练时需要注意同步BNtorch.nn.SyncBatchNorm.convert_sync_batchnorm在单卡上不要启用否则每个step会有一次额外的loss通信速度反而变慢。跑对比实验时两个模型放在同一个batch size、同一种数据增强、同一个优化器设置下才有意义否则你对比的是炼丹配方而不是网络结构。5. 训练脚本、评估指标与UNet和DeepLabV3的推理调优5.1 损失函数、优化器与学习率策略选择语义分割的训练脚本一般共用下面这套配置模板criterion nn.CrossEntropyLoss(ignore_index255) optimizer torch.optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-4) scheduler torch.optim.lr_scheduler.PolyLR(optimizer, total_iters10000, power0.9) for epoch in range(30): model.train() for images, masks in dataloader: images images.cuda() masks masks.cuda() outputs model(images) # 注意类型outputs是floatmasks是long loss criterion(outputs, masks) optimizer.zero_grad() loss.backward() optimizer.step() scheduler.step()ignore_index255表示标签里255像素位置不参与loss计算这专门用来处理VOC数据集里边框上的难例区域。PolyLR是语义分割里面的主流选择它每步按(1 - iter/total_iters)^power衰减比StepLR细致。AdamW配合weight_decay一般适合分割SGDmomentum如果初始学习率设成0.007效果也相近但训练曲线更陡小batch不太好压住。UNet训练自己的数据集时类别不平衡很常见。一张512x512的车道线图背景往往占95%以上直接用CrossEntropy会让模型把所有像素预测成背景。这时可以给criterion传weight参数把前景类权重设成类别频率反比或者换DiceBCELoss混合损失。更简单粗暴的是在Dataset里对前景样本做一番重采样让每张batch里至少包含一定比例的前景区域。5.2 评估指标的实现mIoU和Pixel Accuracy训练过程看loss评估模型好坏看mIoU。mIoU的计算逻辑是对每个类别求预测正确且标签正确的像素数除以该类别预测区域和真实区域的并集面积。def compute_miou(pred, mask, num_classes21): pred是网络输出的logitsshape(N,C,H,W)mask是标签(N,H,W) pred pred.argmax(dim1) # 取概率最大类别 iou_list [] pred pred.cpu().numpy() mask mask.cpu().numpy() for cls in range(num_classes): intersection ((pred cls) (mask cls)).sum() union ((pred cls) | (mask cls)).sum() if union 0: iou_list.append(float(nan)) # 类别不存在时跳过 else: iou_list.append(intersection / union) # 忽略nan算所有类别的平均 valid_iou [iou for iou in iou_list if not np.isnan(iou)] return np.mean(valid_iou)这个实现里隐藏着一个性能问题循环num_classes次全图比较VOC的21类在512x512图上单次评估就要约1秒整个验证集要跑几分钟。工程上会改成矩阵运算加速一次性算全部类别的混淆矩阵。评估时记得把model.eval()和torch.no_grad()同时用上前者关掉Dropout和BN更新后者关掉梯度图构建两者都不可少。5.3 几个能直接提升推理效果的落地技巧换模型参数之前先检查输入预处理是否和训练时完全一致。很多UNet项目训练时用的是Normalize后的Tensor推理时却直接传了PIL图片或BGR格式精度直接崩掉。推理封装里把ToTensor和Normalize按训练时参数同步复制一份。预测结果的原始输出是(H, W)的类别索引图可视化时用调色板映射成彩色图。最常见做法是预先定义21维的RGB数组索引直接查表着色然后与原图0.5权重做addWeighted混合输出。这样叠加出来的效果图便于检查模型预测边界和错分区域车内做展示时也直观。后处理的最有效手段是TTA多尺度推理把输入缩放到0.75、1.0、1.25倍分别预测把不同尺度的概率图插值回原始分辨率后取平均再取argmax。这个方法对UNet和DeepLabV3都有1-3个点的mIoU提升代价是推理时间翻三倍。另一个是CRF后处理用条件随机场对边缘做平滑修正对UNet的边缘锯齿有一定改善但Python版的pydensecrf库在Windows下安装容易出问题服务器上用之前先跑一遍pip install pydensecrf验证通过再写进推理管线。最后值得单独测试的是混合精度推理。torch.autocast(device_typecuda, dtypetorch.float16)配合torch.inference_mode()在T4、V100这种GPU上能把单帧推理时间压缩30%左右分割精度损失通常不超过0.5个点。跑通这一步整个项目从训练、评估到部署的链路就完整了。本文还有配套的精品资源点击获取
RELATED — 相关阅读

相关资讯

LATEST — 最新资讯

最新发布

TODAY — 本日精选

新闻

WEEKLY — 本周精选

新闻

MONTHLY — 本月精选

新闻