尧图网站建设 
简介这是一份基于Pytorch实现DeepLabv3图像分割算法的实战项目面向有一定深度学习基础、希望掌握语义分割模型训练的读者。资源在VOC与Cityscapes数据集上完成训练与预测覆盖从数据加载、网络构建、损失函数到评估可视化的完整流程。压缩包共55个文件大小仅2.25MB主体为23个Python脚本含训练主程序、预测脚本、网络定义及工具函数并配有17张示例结果图和2个txt数据划分文件另有README详细流程教程及备份文件目录结构清晰方便对照学习。已有139人学习下载。通过该项目读者可以直观理解DeepLabv3的代码实现学会在自定义数据上训练模型并利用可视化模块快速检验分割效果从数据预处理到模型评估均有对应脚本是一份适合入门与复现的轻量级参考资料。1. 图像分割任务与DeepLabv3要解决的两个痛点做分割训练时模型结构往往不是第一生产力数据才是。第一次在Pytorch上跑DeepLabv3的人有一半时间会花在报错上VOC的调色板PNG读出来是四通道、Cityscapes的trainId和evalId对不上、损失下降到某个数值后mIoU就卡住不动。这些坑说穿了都是数据和标签格式的问题和模型本身关系不大。DeepLabv3是语义分割里最常被拿来做基线的模型它用空洞卷积扩大感受野再用Decoder恢复边缘细节在VOC和Cityscapes上都能稳定跑到不错的精度。这篇文章从环境搭建开始把数据管线、ASPP和Decoder的Pytorch实现、训练循环、mIoU评估这几个环节依次讲透适合有分类模型训练经验、但第一次做像素级标注任务的人。2. Pytorch环境准备与VOC/Cityscapes数据管线的搭建2.1 安装Pytorch的CUDA版本匹配与验证Pytorch安装教程里最常见的问题不是装不上而是装完之后torch.cuda.is_available()返回False。原因通常是conda装了CPU版本或者CUDA toolkit版本和驱动不匹配。先分开看两件事nvidia-smi显示的Driver Version决定你最多能用哪个CUDA版本而Pytorch安装时选的pytorch-cuda是运行时依赖的CUDA runtime两者不是一个东西。驱动足够新如525以上就能支持12.1的Pytorch包。# 查看驱动支持的CUDA版本 nvidia-smi # 创建独立环境避免污染基础环境 conda create -n deeplab python3.10 -y conda activate deeplab # 安装CUDA 12.1对应的Pytorch 2.1.0 conda install pytorch2.1.0 torchvision0.16.0 pytorch-cuda12.1 -c pytorch -c nvidia验证是否安装成功用一小段代码跑一个矩阵乘法和设备检测import torch print(torch.__version__) print(torch.cuda.is_available()) x torch.rand(1024, 1024).cuda() y torch.rand(1024, 1024).cuda() print((x y).shape)参数说明python3.10是Pytorch官方支持较好的版本TensorFlow和Pytorch两套环境互不干扰pytorch-cuda12.1表示使用CUDA 12.1的预编译包如果机器驱动是11.x就换成11.8。不要手动装CUDA toolkit再装CPU版Pytorch这样最容易出现版本错乱。如果你用的是Ubuntu系统流程完全一致conda会帮你处理底层依赖。2.2 数据集目录与标注格式VOC的调色板PNG与Cityscapes的trainIdVOC和Cityscapes是两种完全不同的标注体系。VOC的SegmentationClass目录下每张PNG是带调色板的单通道索引图直接用PIL.Image.open()读出来是P模式需要转成RGB或者转成numpy array以后才能看到真实类别ID很多人在这里踩坑。Cityscapes的gtFine目录下同时提供了labelIds和labelTrainIds两套PNG后者才是训练时要用的它的像素值已经映射成0到18的类别ID255留给忽略区域。常见的目录结构如下VOC2012/ ├── JPEGImages/ ├── SegmentationClass/ └── ImageSets/ └── Segmentation/ ├── train.txt └── val.txt Cityscapes/ ├── leftImg8bit/ │ ├── train/ │ └── val/ ├── gtFine/ │ ├── train/ │ └── val/ └── meta/ └── class_mapping.csv两张数据集的差异对照如下项目VOC2012Cityscapes类别数21含背景19忽略区域用255标注格式调色板索引PNGlabelTrainIds.png原图尺寸不等约500x3752048x1024训练规模约10k张约3k张难点类别不平衡小目标多、边缘精细Cityscapes官方原图是2048x1024直接送进模型显存很可能不够常见做法是随机裁剪成769x769或者512x512。VOC因为本身是小图一般直接resize到512x512训练起来速度快很多。2.3 自定义SegDataset和DataLoader的参数细节Pytorch的torchvision.datasets.VOCSegmentation可以直接读取VOC但Cityscapes没有内置在torchvision里。自己写Dataset要同时兼容两套路径和两种标注格式下面的写法把两套数据统一成一个接口import os import numpy as np from PIL import Image from torch.utils.data import Dataset class SegDataset(Dataset): def __init__(self, image_dir, mask_dir, data_list, base_size(512, 512), is_cityscapesTrue): self.image_paths [] self.mask_paths [] for line in open(data_list, r): line line.strip() self.image_paths.append(os.path.join(image_dir, line .jpg)) self.mask_paths.append(os.path.join(mask_dir, line .png)) self.base_size base_size self.is_cityscapes is_cityscapes def __len__(self): return len(self.image_paths) def __getitem__(self, idx): img Image.open(self.image_paths[idx]).convert(RGB) mask Image.open(self.mask_paths[idx]) if self.is_cityscapes: # Cityscapes的mask已经是trainId索引直接转数组 mask np.array(mask, dtypenp.int64) else: # VOC的mask是调色板索引np.array后本身就是0~20的索引 mask np.array(mask, dtypenp.int64) img img.resize(self.base_size, Image.BILINEAR) mask Image.fromarray(mask).resize( self.base_size, Image.NEAREST) mask np.array(mask, dtypenp.int64) img np.array(img, dtypenp.float32) img img.transpose(2, 0, 1) # HWC - CHW img torch.from_numpy(img) / 255.0 mask torch.from_numpy(mask) return img, mask关键参数说明is_cityscapesTrue时不需要任何转换直接读像素值因为Cityscapes的labelTrainIds.png已经是训练IDVOC的PNG虽然也是索引图但颜色映射和类别ID是绑定在文件本身的np.array()得到的值就是0到20。resize时mask必须用Image.NEAREST否则会平滑出无意义的中间值这一点对任何分割任务都适用。训练时不要在原图尺寸上直接做全图训练建议配合随机裁剪增强。DataLoader侧的参数也值得花时间调num_workers在Linux上可以开到CPU核数的一半Windows上建议不超过4否则容易报DataLoader worker error。pin_memoryTrue能让GPU拷贝快一些前提是数据在CUDA页锁定内存上。from torch.utils.data import DataLoader train_loader DataLoader( train_dataset, batch_size8, shuffleTrue, num_workers6, pin_memoryTrue, drop_lastTrue )提示drop_lastTrue在分割任务里尽量保持因为Cityscapes的batch只有3k张最后一批不够batch_size时BatchNorm的统计量会偏差很大。3. 用Pytorch复现DeepLabv3的ASPP与Decoder结构3.1 从ResNet到空洞卷积输出为什么是stride16DeepLab系列的核心思想是在不降低特征图分辨率的前提下扩大感受野。普通ResNet最后一层输出是输入图像的1/32对分割来说丢失了太多边界信息。DeepLabv3把ResNet的layer3和layer4的stride改成1用空洞卷积来补偿下采样带来的感受野损失最终特征图是输入的1/16也就是常说的output stride16OS16。这样做让特征图分辨率翻了一倍同时感受野并没有缩小。选择OS16而不是OS8作为默认值是因为OS8需要把layer4的stride也改成1ResNet最后两个stage都变成空洞模式显存占用显著上升。大部分训练场景下OS16配合Decoder已经能取得很好的精度OS8更多用在推理阶段做精度提升。3.2 ASPP多尺度特征提取的Pytorch实现ASPPAtrous Spatial Pyramid Pooling是DeepLabv3的核心模块它把同一张特征图分别用不同空洞率的卷积并行处理再拼接起来。空洞率rate决定了卷积核的采样间隔rate越大看到的范围越广但取得更稀疏。Pytorch的nn.Conv2d直接支持dilation参数不需要额外写采样逻辑。import torch import torch.nn as nn import torch.nn.functional as F class ASPP(nn.Module): def __init__(self, in_channels2048, out_channels256, rates(6, 12, 18)): super(ASPP, self).__init__() self.branches nn.ModuleList() # 1x1卷积分支保留原始尺度的语义信息 self.branches.append(nn.Sequential( nn.Conv2d(in_channels, out_channels, 1, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) )) # 三个不同空洞率的3x3空洞卷积 for rate in rates: self.branches.append(nn.Sequential( nn.Conv2d(in_channels, out_channels, 3, paddingrate, dilationrate, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) )) # 全局平均池化分支获取整图的全局上下文 self.gap_branch nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Conv2d(in_channels, out_channels, 1, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) ) # 五个分支拼接后投影回out_channels self.project nn.Sequential( nn.Conv2d(out_channels * 5, out_channels, 1, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue), nn.Dropout(0.1) ) def forward(self, x): size x.shape[-2:] branch_out [branch(x) for branch in self.branches] # 全局池化分支需要上采样回原尺寸再拼接 gap self.gap_branch(x) gap F.interpolate(gap, sizesize, modebilinear, align_cornersTrue) branch_out.append(gap) return self.project(torch.cat(branch_out, dim1))这里paddingrate和dilationrate必须保持一致保证卷积输出尺寸不变。之所以不使用空洞率24是因为rate过大会导致卷积核采样的有效像素过少效果反而下降这也是后续很多改进版本把rates改成(1, 2, 4)连续级别的原因。Dropout(0.1)是固定的在主干网络后面放很小的dropout能缓解过拟合不会像大dropout那样让模型欠拟合。Pytorch官方在torchvision里也提供了torchvision.models.segmentation.deeplabv3_resnet101但那是DeepLabv3而不是v3结构上少了Decoder模块。直接用预训练权重可以结构需要自己补上Decoder才能对齐原论文。3.3 Decoder整合低层细节的实现细节DeepLabv3与DeepLabv3最大的差别就是增加了一个轻量Decoder。ASPP输出的特征已经下采样了16倍直接双线性上采样回原图会丢失边缘细节Decoder把ResNet第一个stage的低层特征下采样4倍的那个也接入进来通过1x1卷积降维后和ASPP输出拼接。低层特征通道数少、空间信息多高层特征通道数多、语义信息强两者互补。class Decoder(nn.Module): def __init__(self, low_level_channels256, num_classes21): super(Decoder, self).__init__() # 低层特征先降到48通道减少参数量同时抑制噪声 self.low_level_conv nn.Sequential( nn.Conv2d(low_level_channels, 48, 1, biasFalse), nn.BatchNorm2d(48), nn.ReLU(inplaceTrue) ) # 两个3x3卷积融合最后接分类卷积 self.conv nn.Sequential( nn.Conv2d(256 48, 256, 3, padding1, biasFalse), nn.BatchNorm2d(256), nn.ReLU(inplaceTrue), nn.Conv2d(256, 256, 3, padding1, biasFalse), nn.BatchNorm2d(256), nn.ReLU(inplaceTrue) ) self.classifier nn.Conv2d(256, num_classes, 1) def forward(self, high_level_feat, low_level_feat): # 高维特征先上采样到低层特征相同的尺寸 high_level_feat F.interpolate( high_level_feat, sizelow_level_feat.shape[-2:], modebilinear, align_cornersTrue ) low self.low_level_conv(low_level_feat) # 拼接后过3x3卷积最后输出各类别score map x torch.cat([low, high_level_feat], dim1) x self.conv(x) return self.classifier(x)拼接维度是256加48等于304这是DeepLabv3结构里经常写错的地方。low_level_channels要根据backbone决定ResNet101的layer1输出是256通道所以是256如果换用MobileNetV2作为backbonelayer1输出是24通道相应的48也要改成24附近的值。最后的1x1卷积就是逐像素分类器输出维度是类别数VOC为21Cityscapes为19。4. 训练DeepLabv3损失函数、学习率策略与训练循环写法4.1 损失函数选择与类别不平衡分割任务默认用nn.CrossEntropyLoss但它对长尾分布不友好。VOC里背景像素占了一半以上Cityscapes的细长物体如交通灯、摩托车更是极少数像素。有三种常见处理方式给少数类加大权重、用Focal Loss、或者在训练过程中做Online Hard Example MiningOHEM。实际工程里最稳妥的还是给CrossEntropyLoss传一个weight向量权重根据各类别像素占比的倒数来设定。def compute_class_weights(mask_dir, num_classes): 统计所有mask中每个类别的像素数量 权重设为总像素数 / (类别数 * 该类像素数) counts np.zeros(num_classes, dtypenp.float64) for mask_path in mask_dir: mask np.array(Image.open(mask_path)) for cls in range(num_classes): counts[cls] (mask cls).sum() total counts.sum() weights total / (num_classes * counts 1e-6) # 抑制极端权重防止小类权重过大导致震荡 weights np.clip(weights, 0.1, 10.0) return torch.tensor(weights, dtypetorch.float32) criterion nn.CrossEntropyLoss( weightcompute_class_weights(...).cuda(), ignore_index255 )参数说明keep统计前确认Cityscapes不需要统计255忽略区域ignore_index255告诉损失函数那些空白标注不参与梯度计算clip到0.1到10之间很关键VOC的“猫”这类类别像素极少不限制权重会达到几十倍训练初期直接崩掉。Focal Loss在分割任务上提升有限更多用在目标检测里建议先跑通常规CE再考虑替换。4.2 poly学习率与WarmUp的Pytorch实现语义分割和分类任务有一个显著区别分类模型的step学习率配合cosine能跑得很好分割任务里最经典的是poly学习率策略公式是lr base_lr * (1 - iter/total_iter) ^ powerpower通常取0.9。这种策略在训练后期让学习率慢慢降到接近零能显著提升收敛精度。直接在每个step里手动修改学习率比用torch.optim.lr_scheduler更容易控制。def poly_lr(optimizer, base_lr, cur_iter, total_iters, power0.9): lr base_lr * (1 - cur_iter / total_iters) ** power for param_group in optimizer.param_groups: param_group[lr] lr return lr # 训练循环里每个step调用一次 for i, (img, mask) in enumerate(train_loader): cur_lr poly_lr(optimizer, 0.01, epoch * len(train_loader) i, num_epochs * len(train_loader)) optimizer.zero_grad() ...WarmUp一般加在训练的前几轮做法是让学习率从base_lr的十分之一线性升到base_lr。DeepLabv3在ImageNet预训练权重下不需要过长warmup500步以内足够。如果是从头训练warmup对最终精度影响很大建议warmup到5个epoch。4.3 训练主循环与验证集mIoU监控一次完整的训练循环下面这段代码可以作为骨架验证集评估函数稍后单独讲训练时每隔几个epoch调用一次根据mIoU保存最优权重。def train_one_epoch(model, loader, optimizer, criterion, cur_epoch, total_epochs, total_iters): model.train() running_loss 0.0 for i, (images, masks) in enumerate(loader): images images.cuda() masks masks.cuda() base_lr 0.01 cur_lr poly_lr(optimizer, base_lr, cur_epoch * len(loader) i, total_iters) optimizer.zero_grad() outputs model(images)[out] # 或者直接outputs loss criterion(outputs, masks) loss.backward() optimizer.step() running_loss loss.item() if i % 50 0: print(fEpoch {cur_epoch}, Batch {i}, Loss {loss.item():.4f}, fLr {cur_lr:.6f}) return running_loss / len(loader)注意model输出的结构不一定都是tensor用torchvision.models.segmentation里的模型返回的是dict需要取[out]自己搭的模型返回的是原样tensor写代码时要统一。每个batch结束后手动调用poly_lr而不是依赖scheduler这样能精确到iter级别调整。验证放在训练epoch和epoch之间Cityscapes每轮验证全量val集大约需要5分钟VOC更快一些建议至少每两个epoch验证一次。5. 评估指标与Cityscapes可视化验证5.1 mIoU计算的numpy实现与混淆矩阵mIoUMean Intersection over Union是分割任务里最核心的指标计算方式是对每个类别计算预测区域和真实区域的交并比然后对所有类取平均。直接遍历每个像素计算IoU会很慢标准做法是用混淆矩阵。混淆矩阵第i行第j列表示真实类别是i、预测类别是j的像素数。import numpy as np class SegmentationMetric: def __init__(self, num_classes, ignore_index255): self.num_classes num_classes self.ignore_index ignore_index self.hist np.zeros((num_classes, num_classes), dtypenp.int64) def add_batch(self, pred, label): pred: [B, H, W] 的argmax后的类别索引 label: [B, H, W] 的真实索引可能包含255 pred pred.astype(np.int64) label label.astype(np.int64) mask (label ! self.ignore_index) pred pred[mask] label label[mask] # 用索引编码构造混淆矩阵 idx label * self.num_classes pred hist np.bincount(idx, minlengthself.num_classes ** 2) self.hist hist.reshape(self.num_classes, self.num_classes) def compute_miou(self): inter np.diag(self.hist) union self.hist.sum(axis1) self.hist.sum(axis0) - inter iou inter / np.maximum(union, 1) miou np.mean(iou) return miou, iou说明np.bincount比双重循环快两个数量级Cityscapes一张1024x2048的mask全量遍历大约2万个像素索引minlength必须显式指定union的计算要防止分母为零用np.maximum(union, 1)而不是直接除否则一个完全没预测到的类会出现NaN。返回的iou数组可以按索引对应每个类别的单独表现训练到后期要关注IoU最低的是哪几个类通常是最小的类别这说明模型在小目标上还有空间。5.2 ignore_index与Cityscapes的evalId映射Cityscapes的特殊之处在于训练类别有19个但在评测时某些类别在最终计算指标时要合并或者忽略。官方给了一套trainId到evalId的映射比如trainId7的road在evalId里还是0但一些细分类如trainId17和trainId18摩托车和自行车合并为vehicle。自己写评测时最简单的做法是保持trainId体系直接算19类mIoU这样能反映训练的真实效果要发布对比论文时才需要转成evalId的16类别。转换逻辑如下# trainId - evalId长度19 train_id_to_eval_id [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18] # 如果某个类别在eval中被忽略设为-1禁用 eval_mask np.array([True] * 19) eval_mask[2] False # 假设类别2忽略实际以官方为准Cityscapes的ignore_index255在训练和验证中都必须保持一致否则验证时背景区域会按照类别0参与mIoU计算把指标拉低几个点。对比实验时一定要固定ignore的设定。5.3 预测可视化边缘误差与感受野mIoU只是个数字想要定位模型哪里做得不好必须可视化预测结果。推理时把模型的输出做argmax映射到颜色空间和GT并排对比。VOC有固定的颜色映射表Cityscapes也有官方颜色定义直接用PIL保存为PNG即可。def vis_prediction(model, img_tensor, color_map, save_path): model.eval() with torch.no_grad(): out model(img_tensor.unsqueeze(0).cuda()) if isinstance(out, dict): out out[out] pred out.argmax(dim1).squeeze(0).cpu().numpy() # 将类别索引映射为RGB颜色 h, w pred.shape rgb np.zeros((h, w, 3), dtypenp.uint8) for cls, color in enumerate(color_map): rgb[pred cls] color Image.fromarray(rgb).save(save_path)观察输出的图重点看两个区域一是物体的边界DeepLabv3边界效果差通常出在Decoder的浅层特征提取不够干净低层特征的1x1卷积是直接在原图分辨率上操作的噪音会被放大二是特别小的物体如果小物体整体丢失往往是ASPP的全局池化分支权重过高考虑把ASPP的out_ch从256改成512看是否有改善。6. 提升DeepLabv3精度的几个具体调参思路6.1 加载ImageNet预训练权重的正确姿势分割网络从头训练收敛极慢且效果差必须加载backbone在ImageNet上的预训练权重。ResNet101的权重在torchvision.models里有官方版本加载时要注意模型结构里的分类头要去掉。如果自己定义了带空洞卷积的ResNet层的名字和官方默认版本不同直接load会报strict错误需要把strict设成False再手动检查哪些层的参数没对齐。import torchvision.models as models resnet models.resnet101(pretrainedTrue) backbone_state resnet.state_dict() # 以自己定义的backbone名字为准做键名匹配加载 model.backbone.load_state_dict(backbone_state, strictFalse)提示pretrainedTrue的下载需要保持网络通畅第一次使用先把权重下好存到本地后续训练不用重复下载。如果加载后训练指标波动特别大优先检查BN层的running_mean和running_var有没有被错误初始化成零。6.2 output stride用OS8代替OS16训练时用OS16速度更快inference时切换成OS8来获得更精细的分割结果是DeepLabv3论文里明确提到的策略。实现方式是在加载完整模型后把layer4的stride改成1同时把最后一个stage的3x3卷积的dilation从1改成2或者直接整体降低backbone的stride。这个操作在验证时做可以涨0.5到1个点的mIoU且不额外增加参数量。6.3 小显存训练的batch size替代方案Cityscapes原图768x769的batch size8在12G显存上勉强能跑但如果你只有8G显存用OS16和batch size4会明显掉精度。常见应对是保持batch size8把裁剪尺寸从769改成640DeepLabv3对输入尺寸不算特别敏感640输入在VOC上约损失0.3个点但能让梯度更稳定。更进一步的方案是使用梯度累积每4个batch更新一次效果等同于batch size翻倍accum_steps 4 for i, (img, mask) in enumerate(train_loader): loss criterion(model(img)[out], mask) / accum_steps loss.backward() if (i 1) % accum_steps 0: optimizer.step() optimizer.zero_grad()另外BatchNorm在batch size小于8时统计量不稳定建议把BN的momentum从默认0.1调到0.99让running均值更平稳这样小batch下训练出来的模型在验证集上波动更小。本文还有配套的精品资源点击获取 PyTorch实现DeepLabv3+：语义分割模型完整实战指南
简介这是一份基于Pytorch实现DeepLabv3图像分割算法的实战项目面向有一定深度学习基础、希望掌握语义分割模型训练的读者。资源在VOC与Cityscapes数据集上完成训练与预测覆盖从数据加载、网络构建、损失函数到评估可视化的完整流程。压缩包共55个文件大小仅2.25MB主体为23个Python脚本含训练主程序、预测脚本、网络定义及工具函数并配有17张示例结果图和2个txt数据划分文件另有README详细流程教程及备份文件目录结构清晰方便对照学习。已有139人学习下载。通过该项目读者可以直观理解DeepLabv3的代码实现学会在自定义数据上训练模型并利用可视化模块快速检验分割效果从数据预处理到模型评估均有对应脚本是一份适合入门与复现的轻量级参考资料。1. 图像分割任务与DeepLabv3要解决的两个痛点做分割训练时模型结构往往不是第一生产力数据才是。第一次在Pytorch上跑DeepLabv3的人有一半时间会花在报错上VOC的调色板PNG读出来是四通道、Cityscapes的trainId和evalId对不上、损失下降到某个数值后mIoU就卡住不动。这些坑说穿了都是数据和标签格式的问题和模型本身关系不大。DeepLabv3是语义分割里最常被拿来做基线的模型它用空洞卷积扩大感受野再用Decoder恢复边缘细节在VOC和Cityscapes上都能稳定跑到不错的精度。这篇文章从环境搭建开始把数据管线、ASPP和Decoder的Pytorch实现、训练循环、mIoU评估这几个环节依次讲透适合有分类模型训练经验、但第一次做像素级标注任务的人。2. Pytorch环境准备与VOC/Cityscapes数据管线的搭建2.1 安装Pytorch的CUDA版本匹配与验证Pytorch安装教程里最常见的问题不是装不上而是装完之后torch.cuda.is_available()返回False。原因通常是conda装了CPU版本或者CUDA toolkit版本和驱动不匹配。先分开看两件事nvidia-smi显示的Driver Version决定你最多能用哪个CUDA版本而Pytorch安装时选的pytorch-cuda是运行时依赖的CUDA runtime两者不是一个东西。驱动足够新如525以上就能支持12.1的Pytorch包。# 查看驱动支持的CUDA版本 nvidia-smi # 创建独立环境避免污染基础环境 conda create -n deeplab python3.10 -y conda activate deeplab # 安装CUDA 12.1对应的Pytorch 2.1.0 conda install pytorch2.1.0 torchvision0.16.0 pytorch-cuda12.1 -c pytorch -c nvidia验证是否安装成功用一小段代码跑一个矩阵乘法和设备检测import torch print(torch.__version__) print(torch.cuda.is_available()) x torch.rand(1024, 1024).cuda() y torch.rand(1024, 1024).cuda() print((x y).shape)参数说明python3.10是Pytorch官方支持较好的版本TensorFlow和Pytorch两套环境互不干扰pytorch-cuda12.1表示使用CUDA 12.1的预编译包如果机器驱动是11.x就换成11.8。不要手动装CUDA toolkit再装CPU版Pytorch这样最容易出现版本错乱。如果你用的是Ubuntu系统流程完全一致conda会帮你处理底层依赖。2.2 数据集目录与标注格式VOC的调色板PNG与Cityscapes的trainIdVOC和Cityscapes是两种完全不同的标注体系。VOC的SegmentationClass目录下每张PNG是带调色板的单通道索引图直接用PIL.Image.open()读出来是P模式需要转成RGB或者转成numpy array以后才能看到真实类别ID很多人在这里踩坑。Cityscapes的gtFine目录下同时提供了labelIds和labelTrainIds两套PNG后者才是训练时要用的它的像素值已经映射成0到18的类别ID255留给忽略区域。常见的目录结构如下VOC2012/ ├── JPEGImages/ ├── SegmentationClass/ └── ImageSets/ └── Segmentation/ ├── train.txt └── val.txt Cityscapes/ ├── leftImg8bit/ │ ├── train/ │ └── val/ ├── gtFine/ │ ├── train/ │ └── val/ └── meta/ └── class_mapping.csv两张数据集的差异对照如下项目VOC2012Cityscapes类别数21含背景19忽略区域用255标注格式调色板索引PNGlabelTrainIds.png原图尺寸不等约500x3752048x1024训练规模约10k张约3k张难点类别不平衡小目标多、边缘精细Cityscapes官方原图是2048x1024直接送进模型显存很可能不够常见做法是随机裁剪成769x769或者512x512。VOC因为本身是小图一般直接resize到512x512训练起来速度快很多。2.3 自定义SegDataset和DataLoader的参数细节Pytorch的torchvision.datasets.VOCSegmentation可以直接读取VOC但Cityscapes没有内置在torchvision里。自己写Dataset要同时兼容两套路径和两种标注格式下面的写法把两套数据统一成一个接口import os import numpy as np from PIL import Image from torch.utils.data import Dataset class SegDataset(Dataset): def __init__(self, image_dir, mask_dir, data_list, base_size(512, 512), is_cityscapesTrue): self.image_paths [] self.mask_paths [] for line in open(data_list, r): line line.strip() self.image_paths.append(os.path.join(image_dir, line .jpg)) self.mask_paths.append(os.path.join(mask_dir, line .png)) self.base_size base_size self.is_cityscapes is_cityscapes def __len__(self): return len(self.image_paths) def __getitem__(self, idx): img Image.open(self.image_paths[idx]).convert(RGB) mask Image.open(self.mask_paths[idx]) if self.is_cityscapes: # Cityscapes的mask已经是trainId索引直接转数组 mask np.array(mask, dtypenp.int64) else: # VOC的mask是调色板索引np.array后本身就是0~20的索引 mask np.array(mask, dtypenp.int64) img img.resize(self.base_size, Image.BILINEAR) mask Image.fromarray(mask).resize( self.base_size, Image.NEAREST) mask np.array(mask, dtypenp.int64) img np.array(img, dtypenp.float32) img img.transpose(2, 0, 1) # HWC - CHW img torch.from_numpy(img) / 255.0 mask torch.from_numpy(mask) return img, mask关键参数说明is_cityscapesTrue时不需要任何转换直接读像素值因为Cityscapes的labelTrainIds.png已经是训练IDVOC的PNG虽然也是索引图但颜色映射和类别ID是绑定在文件本身的np.array()得到的值就是0到20。resize时mask必须用Image.NEAREST否则会平滑出无意义的中间值这一点对任何分割任务都适用。训练时不要在原图尺寸上直接做全图训练建议配合随机裁剪增强。DataLoader侧的参数也值得花时间调num_workers在Linux上可以开到CPU核数的一半Windows上建议不超过4否则容易报DataLoader worker error。pin_memoryTrue能让GPU拷贝快一些前提是数据在CUDA页锁定内存上。from torch.utils.data import DataLoader train_loader DataLoader( train_dataset, batch_size8, shuffleTrue, num_workers6, pin_memoryTrue, drop_lastTrue )提示drop_lastTrue在分割任务里尽量保持因为Cityscapes的batch只有3k张最后一批不够batch_size时BatchNorm的统计量会偏差很大。3. 用Pytorch复现DeepLabv3的ASPP与Decoder结构3.1 从ResNet到空洞卷积输出为什么是stride16DeepLab系列的核心思想是在不降低特征图分辨率的前提下扩大感受野。普通ResNet最后一层输出是输入图像的1/32对分割来说丢失了太多边界信息。DeepLabv3把ResNet的layer3和layer4的stride改成1用空洞卷积来补偿下采样带来的感受野损失最终特征图是输入的1/16也就是常说的output stride16OS16。这样做让特征图分辨率翻了一倍同时感受野并没有缩小。选择OS16而不是OS8作为默认值是因为OS8需要把layer4的stride也改成1ResNet最后两个stage都变成空洞模式显存占用显著上升。大部分训练场景下OS16配合Decoder已经能取得很好的精度OS8更多用在推理阶段做精度提升。3.2 ASPP多尺度特征提取的Pytorch实现ASPPAtrous Spatial Pyramid Pooling是DeepLabv3的核心模块它把同一张特征图分别用不同空洞率的卷积并行处理再拼接起来。空洞率rate决定了卷积核的采样间隔rate越大看到的范围越广但取得更稀疏。Pytorch的nn.Conv2d直接支持dilation参数不需要额外写采样逻辑。import torch import torch.nn as nn import torch.nn.functional as F class ASPP(nn.Module): def __init__(self, in_channels2048, out_channels256, rates(6, 12, 18)): super(ASPP, self).__init__() self.branches nn.ModuleList() # 1x1卷积分支保留原始尺度的语义信息 self.branches.append(nn.Sequential( nn.Conv2d(in_channels, out_channels, 1, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) )) # 三个不同空洞率的3x3空洞卷积 for rate in rates: self.branches.append(nn.Sequential( nn.Conv2d(in_channels, out_channels, 3, paddingrate, dilationrate, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) )) # 全局平均池化分支获取整图的全局上下文 self.gap_branch nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Conv2d(in_channels, out_channels, 1, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) ) # 五个分支拼接后投影回out_channels self.project nn.Sequential( nn.Conv2d(out_channels * 5, out_channels, 1, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue), nn.Dropout(0.1) ) def forward(self, x): size x.shape[-2:] branch_out [branch(x) for branch in self.branches] # 全局池化分支需要上采样回原尺寸再拼接 gap self.gap_branch(x) gap F.interpolate(gap, sizesize, modebilinear, align_cornersTrue) branch_out.append(gap) return self.project(torch.cat(branch_out, dim1))这里paddingrate和dilationrate必须保持一致保证卷积输出尺寸不变。之所以不使用空洞率24是因为rate过大会导致卷积核采样的有效像素过少效果反而下降这也是后续很多改进版本把rates改成(1, 2, 4)连续级别的原因。Dropout(0.1)是固定的在主干网络后面放很小的dropout能缓解过拟合不会像大dropout那样让模型欠拟合。Pytorch官方在torchvision里也提供了torchvision.models.segmentation.deeplabv3_resnet101但那是DeepLabv3而不是v3结构上少了Decoder模块。直接用预训练权重可以结构需要自己补上Decoder才能对齐原论文。3.3 Decoder整合低层细节的实现细节DeepLabv3与DeepLabv3最大的差别就是增加了一个轻量Decoder。ASPP输出的特征已经下采样了16倍直接双线性上采样回原图会丢失边缘细节Decoder把ResNet第一个stage的低层特征下采样4倍的那个也接入进来通过1x1卷积降维后和ASPP输出拼接。低层特征通道数少、空间信息多高层特征通道数多、语义信息强两者互补。class Decoder(nn.Module): def __init__(self, low_level_channels256, num_classes21): super(Decoder, self).__init__() # 低层特征先降到48通道减少参数量同时抑制噪声 self.low_level_conv nn.Sequential( nn.Conv2d(low_level_channels, 48, 1, biasFalse), nn.BatchNorm2d(48), nn.ReLU(inplaceTrue) ) # 两个3x3卷积融合最后接分类卷积 self.conv nn.Sequential( nn.Conv2d(256 48, 256, 3, padding1, biasFalse), nn.BatchNorm2d(256), nn.ReLU(inplaceTrue), nn.Conv2d(256, 256, 3, padding1, biasFalse), nn.BatchNorm2d(256), nn.ReLU(inplaceTrue) ) self.classifier nn.Conv2d(256, num_classes, 1) def forward(self, high_level_feat, low_level_feat): # 高维特征先上采样到低层特征相同的尺寸 high_level_feat F.interpolate( high_level_feat, sizelow_level_feat.shape[-2:], modebilinear, align_cornersTrue ) low self.low_level_conv(low_level_feat) # 拼接后过3x3卷积最后输出各类别score map x torch.cat([low, high_level_feat], dim1) x self.conv(x) return self.classifier(x)拼接维度是256加48等于304这是DeepLabv3结构里经常写错的地方。low_level_channels要根据backbone决定ResNet101的layer1输出是256通道所以是256如果换用MobileNetV2作为backbonelayer1输出是24通道相应的48也要改成24附近的值。最后的1x1卷积就是逐像素分类器输出维度是类别数VOC为21Cityscapes为19。4. 训练DeepLabv3损失函数、学习率策略与训练循环写法4.1 损失函数选择与类别不平衡分割任务默认用nn.CrossEntropyLoss但它对长尾分布不友好。VOC里背景像素占了一半以上Cityscapes的细长物体如交通灯、摩托车更是极少数像素。有三种常见处理方式给少数类加大权重、用Focal Loss、或者在训练过程中做Online Hard Example MiningOHEM。实际工程里最稳妥的还是给CrossEntropyLoss传一个weight向量权重根据各类别像素占比的倒数来设定。def compute_class_weights(mask_dir, num_classes): 统计所有mask中每个类别的像素数量 权重设为总像素数 / (类别数 * 该类像素数) counts np.zeros(num_classes, dtypenp.float64) for mask_path in mask_dir: mask np.array(Image.open(mask_path)) for cls in range(num_classes): counts[cls] (mask cls).sum() total counts.sum() weights total / (num_classes * counts 1e-6) # 抑制极端权重防止小类权重过大导致震荡 weights np.clip(weights, 0.1, 10.0) return torch.tensor(weights, dtypetorch.float32) criterion nn.CrossEntropyLoss( weightcompute_class_weights(...).cuda(), ignore_index255 )参数说明keep统计前确认Cityscapes不需要统计255忽略区域ignore_index255告诉损失函数那些空白标注不参与梯度计算clip到0.1到10之间很关键VOC的“猫”这类类别像素极少不限制权重会达到几十倍训练初期直接崩掉。Focal Loss在分割任务上提升有限更多用在目标检测里建议先跑通常规CE再考虑替换。4.2 poly学习率与WarmUp的Pytorch实现语义分割和分类任务有一个显著区别分类模型的step学习率配合cosine能跑得很好分割任务里最经典的是poly学习率策略公式是lr base_lr * (1 - iter/total_iter) ^ powerpower通常取0.9。这种策略在训练后期让学习率慢慢降到接近零能显著提升收敛精度。直接在每个step里手动修改学习率比用torch.optim.lr_scheduler更容易控制。def poly_lr(optimizer, base_lr, cur_iter, total_iters, power0.9): lr base_lr * (1 - cur_iter / total_iters) ** power for param_group in optimizer.param_groups: param_group[lr] lr return lr # 训练循环里每个step调用一次 for i, (img, mask) in enumerate(train_loader): cur_lr poly_lr(optimizer, 0.01, epoch * len(train_loader) i, num_epochs * len(train_loader)) optimizer.zero_grad() ...WarmUp一般加在训练的前几轮做法是让学习率从base_lr的十分之一线性升到base_lr。DeepLabv3在ImageNet预训练权重下不需要过长warmup500步以内足够。如果是从头训练warmup对最终精度影响很大建议warmup到5个epoch。4.3 训练主循环与验证集mIoU监控一次完整的训练循环下面这段代码可以作为骨架验证集评估函数稍后单独讲训练时每隔几个epoch调用一次根据mIoU保存最优权重。def train_one_epoch(model, loader, optimizer, criterion, cur_epoch, total_epochs, total_iters): model.train() running_loss 0.0 for i, (images, masks) in enumerate(loader): images images.cuda() masks masks.cuda() base_lr 0.01 cur_lr poly_lr(optimizer, base_lr, cur_epoch * len(loader) i, total_iters) optimizer.zero_grad() outputs model(images)[out] # 或者直接outputs loss criterion(outputs, masks) loss.backward() optimizer.step() running_loss loss.item() if i % 50 0: print(fEpoch {cur_epoch}, Batch {i}, Loss {loss.item():.4f}, fLr {cur_lr:.6f}) return running_loss / len(loader)注意model输出的结构不一定都是tensor用torchvision.models.segmentation里的模型返回的是dict需要取[out]自己搭的模型返回的是原样tensor写代码时要统一。每个batch结束后手动调用poly_lr而不是依赖scheduler这样能精确到iter级别调整。验证放在训练epoch和epoch之间Cityscapes每轮验证全量val集大约需要5分钟VOC更快一些建议至少每两个epoch验证一次。5. 评估指标与Cityscapes可视化验证5.1 mIoU计算的numpy实现与混淆矩阵mIoUMean Intersection over Union是分割任务里最核心的指标计算方式是对每个类别计算预测区域和真实区域的交并比然后对所有类取平均。直接遍历每个像素计算IoU会很慢标准做法是用混淆矩阵。混淆矩阵第i行第j列表示真实类别是i、预测类别是j的像素数。import numpy as np class SegmentationMetric: def __init__(self, num_classes, ignore_index255): self.num_classes num_classes self.ignore_index ignore_index self.hist np.zeros((num_classes, num_classes), dtypenp.int64) def add_batch(self, pred, label): pred: [B, H, W] 的argmax后的类别索引 label: [B, H, W] 的真实索引可能包含255 pred pred.astype(np.int64) label label.astype(np.int64) mask (label ! self.ignore_index) pred pred[mask] label label[mask] # 用索引编码构造混淆矩阵 idx label * self.num_classes pred hist np.bincount(idx, minlengthself.num_classes ** 2) self.hist hist.reshape(self.num_classes, self.num_classes) def compute_miou(self): inter np.diag(self.hist) union self.hist.sum(axis1) self.hist.sum(axis0) - inter iou inter / np.maximum(union, 1) miou np.mean(iou) return miou, iou说明np.bincount比双重循环快两个数量级Cityscapes一张1024x2048的mask全量遍历大约2万个像素索引minlength必须显式指定union的计算要防止分母为零用np.maximum(union, 1)而不是直接除否则一个完全没预测到的类会出现NaN。返回的iou数组可以按索引对应每个类别的单独表现训练到后期要关注IoU最低的是哪几个类通常是最小的类别这说明模型在小目标上还有空间。5.2 ignore_index与Cityscapes的evalId映射Cityscapes的特殊之处在于训练类别有19个但在评测时某些类别在最终计算指标时要合并或者忽略。官方给了一套trainId到evalId的映射比如trainId7的road在evalId里还是0但一些细分类如trainId17和trainId18摩托车和自行车合并为vehicle。自己写评测时最简单的做法是保持trainId体系直接算19类mIoU这样能反映训练的真实效果要发布对比论文时才需要转成evalId的16类别。转换逻辑如下# trainId - evalId长度19 train_id_to_eval_id [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18] # 如果某个类别在eval中被忽略设为-1禁用 eval_mask np.array([True] * 19) eval_mask[2] False # 假设类别2忽略实际以官方为准Cityscapes的ignore_index255在训练和验证中都必须保持一致否则验证时背景区域会按照类别0参与mIoU计算把指标拉低几个点。对比实验时一定要固定ignore的设定。5.3 预测可视化边缘误差与感受野mIoU只是个数字想要定位模型哪里做得不好必须可视化预测结果。推理时把模型的输出做argmax映射到颜色空间和GT并排对比。VOC有固定的颜色映射表Cityscapes也有官方颜色定义直接用PIL保存为PNG即可。def vis_prediction(model, img_tensor, color_map, save_path): model.eval() with torch.no_grad(): out model(img_tensor.unsqueeze(0).cuda()) if isinstance(out, dict): out out[out] pred out.argmax(dim1).squeeze(0).cpu().numpy() # 将类别索引映射为RGB颜色 h, w pred.shape rgb np.zeros((h, w, 3), dtypenp.uint8) for cls, color in enumerate(color_map): rgb[pred cls] color Image.fromarray(rgb).save(save_path)观察输出的图重点看两个区域一是物体的边界DeepLabv3边界效果差通常出在Decoder的浅层特征提取不够干净低层特征的1x1卷积是直接在原图分辨率上操作的噪音会被放大二是特别小的物体如果小物体整体丢失往往是ASPP的全局池化分支权重过高考虑把ASPP的out_ch从256改成512看是否有改善。6. 提升DeepLabv3精度的几个具体调参思路6.1 加载ImageNet预训练权重的正确姿势分割网络从头训练收敛极慢且效果差必须加载backbone在ImageNet上的预训练权重。ResNet101的权重在torchvision.models里有官方版本加载时要注意模型结构里的分类头要去掉。如果自己定义了带空洞卷积的ResNet层的名字和官方默认版本不同直接load会报strict错误需要把strict设成False再手动检查哪些层的参数没对齐。import torchvision.models as models resnet models.resnet101(pretrainedTrue) backbone_state resnet.state_dict() # 以自己定义的backbone名字为准做键名匹配加载 model.backbone.load_state_dict(backbone_state, strictFalse)提示pretrainedTrue的下载需要保持网络通畅第一次使用先把权重下好存到本地后续训练不用重复下载。如果加载后训练指标波动特别大优先检查BN层的running_mean和running_var有没有被错误初始化成零。6.2 output stride用OS8代替OS16训练时用OS16速度更快inference时切换成OS8来获得更精细的分割结果是DeepLabv3论文里明确提到的策略。实现方式是在加载完整模型后把layer4的stride改成1同时把最后一个stage的3x3卷积的dilation从1改成2或者直接整体降低backbone的stride。这个操作在验证时做可以涨0.5到1个点的mIoU且不额外增加参数量。6.3 小显存训练的batch size替代方案Cityscapes原图768x769的batch size8在12G显存上勉强能跑但如果你只有8G显存用OS16和batch size4会明显掉精度。常见应对是保持batch size8把裁剪尺寸从769改成640DeepLabv3对输入尺寸不算特别敏感640输入在VOC上约损失0.3个点但能让梯度更稳定。更进一步的方案是使用梯度累积每4个batch更新一次效果等同于batch size翻倍accum_steps 4 for i, (img, mask) in enumerate(train_loader): loss criterion(model(img)[out], mask) / accum_steps loss.backward() if (i 1) % accum_steps 0: optimizer.step() optimizer.zero_grad()另外BatchNorm在batch size小于8时统计量不稳定建议把BN的momentum从默认0.1调到0.99让running均值更平稳这样小batch下训练出来的模型在验证集上波动更小。本文还有配套的精品资源点击获取