简介这是一个基于PyTorch实现的LiteSeg实时轻量级语义分割算法资源包面向嵌入式、移动端及自动驾驶等对速度与精度有平衡需求的开发者可用于路面识别、目标检测、医学影像分析等场景。资源共39个文件压缩包约21.09MB其中Python脚本负责模型构建、训练与评估YAML/YML文件提供训练与运行配置JPG/PNG图片为示例与分割效果展示PTH文件为训练好的模型权重DOCX为运行说明文档整体结构清晰便于快速理解项目并开展二次开发。目前已有1358人学习下载。通过这套资源读者可以获得从数据准备、模型训练到部署的全流程参考包括MobileNet/ShuffleNet等轻量骨干网络实现、深度可分离卷积与金字塔特征融合等模块代码、数据增强与损失函数定义以及预训练权重和可视化demo脚本能够直接体验LiteSeg在边缘保持与实时分割上的表现。1. LiteSeg 的实时语义分割定位与架构选择提到实时语义分割很多人第一反应是换一个更小的 Backbone再把输入分辨率降一档但这样省下的计算量往往被精度的断崖式下跌抵消。LiteSeg 的思路是另一条路编码器用 MobileNetV2 这类轻量主干解码器用去掉了大卷积核、把标准卷积替换成深度可分离卷积和空洞卷积组合的 LR-ASPP在几乎不增加参数量的前提下把多尺度上下文补回来。这种结构配合 PyTorch 的动态图和自动混合精度批量不大也能在消费级 GPU 上完成训练。这篇文章面向被显存和算力卡住的工程师和研究同学按训练一套语义分割模型的完整路径来讲覆盖环境搭建、数据集制作、网络实现、损失函数与学习率设计、部署测量你可以直接照着一套 VOC 或 Cityscapes 的实验流程走下来。2. PyTorch 环境搭建与语义分割数据集制作训练 LiteSeg 前的准备2.1 用 Anaconda 建一个干净的 PyTorch 训练环境LiteSeg 本身不要求很新的 PyTorch但训练代码涉及模型定义、数据加载、自动混合精度建议直接用 PyTorch 2.x 配合 CUDA 12.x。我习惯用 Anaconda 隔离环境避免系统 Python 被其他项目的依赖污染。conda create -n liteseg python3.10 -y conda activate liteseg pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121第一条命令创建 Python 3.10 的独立环境第二条激活它。最后一条从 PyTorch 官方源安装 CUDA 12.1 对应的 torch 和 torchvision。这里把 CUDA 版本写进安装索引而不是事后单独装 CUDA是因为 PyTorch 的 wheel 里已经内置了运行时所需的 CUDA 库单独装系统级 CUDA 反而容易出现版本不匹配。验证环境是否可用python -c import torch; print(torch.__version__, torch.cuda.is_available())如果输出2.8.0 True说明 GPU 版本安装成功。只做 CPU 推理的话把安装命令换成pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu。官方源在国外下载较慢可以在命令末尾追加-i https://pypi.tuna.tsinghua.edu.cn/simple使用国内镜像这样 PyTorch 相关依赖包的下载速度会显著提升。提示先确认显卡驱动支持的最高 CUDA 版本再决定安装哪一档 cu118、cu121 还是 cu124驱动版本过低会直接报CUDA driver version is insufficient。2.2 语义分割数据集的目录格式与自定义数据集制作语义分割数据集与分类任务最大的区别在于标注是一张与原始图像相同尺寸的 PNG 图不需要做成 XML 或 JSON。推荐目录组织方式如下data/ images/ train/0001.jpg val/0001.jpg masks/ train/0001.png val/0001.png掩码图的像素值就是类别索引例如背景为 0、人 为 1、车为 2。VOC 数据集的标注里通常有 255 这个值表示该像素不参与损失计算这部分在 Dataset 里通过ignore_index255处理后面训练会用到。制作自定义数据集时常见做法是用 LabelMe 画多边形导出的 JSON 里记录了每个多边形的边点和类别。转成训练掩码的流程分三步先建一张全 0 的单通道图再按多边形填充轮廓内部把对应像素值改成类别索引最后保存成 PNG。注意不要用 JPEG 存掩码JPEG 有损压缩会改变边缘像素值训练时会产生错误标签。PyTorch 侧的数据集类可以这样写from torch.utils.data import Dataset from PIL import Image import numpy as np import torch import os class SegDataset(Dataset): def __init__(self, image_dir, mask_dir, transformNone): self.image_paths sorted( [os.path.join(image_dir, f) for f in os.listdir(image_dir)] ) self.mask_paths sorted( [os.path.join(mask_dir, f) for f in os.listdir(mask_dir)] ) self.transform transform def __len__(self): return len(self.image_paths) def __getitem__(self, idx): image Image.open(self.image_paths[idx]).convert(RGB) mask Image.open(self.mask_paths[idx]) if self.transform: image self.transform(image) mask torch.from_numpy(np.array(mask, dtypenp.int64)) return image, mask这里的核心点是掩码不能走torchvision.transforms.ToTensor()因为ToTensor()会把像素值归一化到 0 到 1导致标签变成浮点数而CrossEntropyLoss要求标签是整数张量。直接np.array(mask, dtypenp.int64)转成 long 型即可像素值 0 到 N-1 正好对应类别索引。上述代码假设图像和掩码文件名一一对应且排序一致实际项目中文件名不一定对齐更稳妥的做法是从同一个文件名列表构造两条路径。提示使用 albumentations 库做数据增强时它会自动同步 image 和 mask 的几何变换不需要像上面这样分别维护两套 transform。但要显式指定mask插值方式为最近邻避免旋转和缩放时标签出现插值产生的虚假类别。3. LiteSeg 的 PyTorch 实现拆解从轻量主干到多尺度解码3.1 MobileNetV2 编码器可换的轻量主干LiteSeg 把 MobileNetV2 当作默认编码器原因是倒残差结构能用更少的 FLOPs 换回相近的表征能力。MobileNetV2 的核心是 InvertedResidual 块先 1x1 卷积升维再 3x3 深度卷积提特征最后 1x1 卷积降维。和传统残差块先压缩再扩展不同这种结构把信息集中在高维空间处理通道间的冗余更低。import torch import torch.nn as nn def conv_bn_relu6(inp, oup, stride1): return nn.Sequential( nn.Conv2d(inp, oup, 3, stride, 1, biasFalse), nn.BatchNorm2d(oup), nn.ReLU6(inplaceTrue) ) class InvertedResidual(nn.Module): def __init__(self, inp, oup, stride, expand_ratio): super().__init__() hidden_dim round(inp * expand_ratio) self.use_res_connect stride 1 and inp oup layers [] if expand_ratio ! 1: layers.append(conv_bn_relu6(inp, hidden_dim, stride1)) layers.extend([ nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groupshidden_dim, biasFalse), nn.BatchNorm2d(hidden_dim), nn.ReLU6(inplaceTrue), nn.Conv2d(hidden_dim, oup, 1, 1, 0, biasFalse), nn.BatchNorm2d(oup), ]) self.conv nn.Sequential(*layers) def forward(self, x): if self.use_res_connect: return x self.conv(x) return self.conv(x)groupshidden_dim就是深度可分离卷积里的 depthwise 部分每个通道单独做卷积参数量从hidden_dim * 9降到hidden_dim * 1。expand_ratio一般取 6也就是中间维度是输入通道的 6 倍。stride为 2 时下采样并且不启用残差连接因为特征图尺寸变了不能直接相加。组装 MobileNetV2 时需要打一个表按输入分辨率递减的顺序堆叠各阶段。通常在最后一个 block 输出的特征图尺寸是输入图像的 1/16 或 1/32对应output_stride16或 32。LiteSeg 希望保留稍高的分辨率来挽救小目标所以常见做法是不在最后一个阶段做 stride2 的下采样改用它后面的空洞卷积保持感受野。实际实现时可以把原本 stride2 的 block 改成 stride1再把同一阶段里所有 3x3 depthwise 卷积设置 dilation2这样特征图停在 1/16 分辨率。3.2 LR-ASPPLiteSeg 的多尺度语义聚合ASPP 的设计思路是在同一个特征图上用不同空洞率的并行空洞卷积采样空洞率越大感受野越大从而同时捕获小目标和大物体。标准 DeepLabV3 里的 ASPP 使用空洞率 6、12、18 的 3x3 卷积外加一个全局平均池化分支在 Cityscapes 上效果不错但计算量偏大。LiteSeg 把它改成 LR-ASPP核心改动是去掉 256 通道的大卷积每个空洞分支只输出 128 通道并且把 3x3 卷积换成深度可分离卷积。class LRASPP(nn.Module): def __init__(self, in_channels, out_channels128, atrous_rates(6, 12, 18)): super().__init__() self.branches nn.ModuleList() for rate in atrous_rates: self.branches.append(nn.Sequential( nn.Conv2d(in_channels, out_channels, 3, 1, paddingrate, dilationrate, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) )) self.image_pool nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Conv2d(in_channels, out_channels, 1, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) ) self.fuse nn.Sequential( nn.Conv2d(out_channels * (len(atrous_rates) 1), out_channels, 1, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) ) def forward(self, x): branch_outs [branch(x) for branch in self.branches] pool_out self.image_pool(x) pool_out nn.functional.interpolate( pool_out, sizex.shape[2:], modebilinear, align_cornersFalse) branch_outs.append(pool_out) return self.fuse(torch.cat(branch_outs, dim1))dilationrate决定了卷积核采样的间距rate 为 6 时 3x3 卷积的实际感受野等效 13x13rate 为 18 时等效 37x37。使用paddingrate保证输出特征图尺寸不变方便后续 concat。全局平均池化分支把整图信息压缩成 1x1再上采样回原尺寸能补全空洞卷积在远距离依赖上的不足。fuse用 1x1 卷积把四个分支拼接后的 512 通道压缩回 128 通道控制解码器计算量。注意空洞卷积是 LiteSeg 这类轻量模型的核心机制但空洞率不是越大越好。rate 超过 24 后卷积核的有效权重会稀疏到几乎只采样几个点容易变成网格伪影实际调参时优先在 6、12、18 附近搜索。3.3 解码器与完整 LiteSeg 网络解码器的常见结构是三层上采样加跳跃连接每次上采样使用双线性插值把特征图分辨率放大 2 倍再与编码器同分辨率的浅层特征相加或拼接。浅层特征包含更多边缘细节深层特征包含更多语义信息跳跃连接就是为了把这两者合并。class LiteSeg(nn.Module): def __init__(self, backbone, aspp, num_classes): super().__init__() self.backbone backbone self.aspp aspp self.decoder nn.Sequential( nn.Conv2d(128 24, 64, 3, 1, 1, biasFalse), nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.Conv2d(64, num_classes, 1) ) def forward(self, x): low_level_feat, high_level_feat self.backbone(x) high_level_feat self.aspp(high_level_feat) high_level_feat nn.functional.interpolate( high_level_feat, sizelow_level_feat.shape[2:], modebilinear, align_cornersFalse) x torch.cat([high_level_feat, low_level_feat], dim1) return self.decoder(x)这里的low_level_feat来自 MobileNetV2 的浅层分辨率是输入的 1/4high_level_feat来自深层分辨率是 1/16。interpolate把深层的 1/16 特征上采样到 1/4然后与浅层特征拼接。decoder里第一个 3x3 卷积的输入通道是128 low_level_feat.channels要根据实际 backbone 输出调整。如果浅层特征取 24 通道拼接后就是 152 通道。如果你不想手写 backbonetorchvision.models里有官方 MobileNetV2可以加载在 ImageNet 上预训练的权重做迁移学习。LiteSeg 训练时微调编码器权重能明显加速收敛比从头训练节省一半以上时间。3.4 验证轻量级模型参数和计算量的度量构建完网络后先量化指标确认它真的符合实时轻量级的定位。用 thop 库统计参数量和 FLOPspip install thopfrom thop import profile, clever_format model LiteSeg(backbone, aspp, num_classes21).cuda() model.eval() flops, params profile(model, inputs(torch.randn(1, 3, 512, 512).cuda(),)) flops, params clever_format([flops, params], %.3f) print(fFLOPs: {flops}, Params: {params})FLOPs 表示模型处理一张 512x512 图像所需的浮点运算次数Params 是模型存储参数量。LiteSeg 在 MobileNetV2 backbone、输出 stride 16 的情况下参数量大约在 4M 到 8M 之间比 DeepLabV3 的 40M 到 60M 小一个数量级。统计时注意 FLOPs 通常不包含 BN 和激活函数的计算不同工具计算口径有差异所以我一般只拿同一工具对比不同模型的相对大小而不是跟论文里数字绝对比较。4. 训练与调参损失函数、学习率与显存策略4.1 损失函数与类别不平衡语义分割最常用的损失是逐像素交叉熵但直接用它处理 Cityscapes 这类类别分布极不均匀的数据集时模型会偏向出现频率高的类别。LiteSeg 的常见做法是交叉熵加 Dice LossDice Loss 直接优化预测区域和真实标注区域的重合度对像素数量少的类别更敏感。class DiceLoss(nn.Module): def __init__(self, smooth1.0, ignore_index255): super().__init__() self.smooth smooth self.ignore_index ignore_index def forward(self, logits, targets): num_classes logits.shape[1] probs torch.softmax(logits, dim1) mask targets ! self.ignore_index targets targets.clone() targets[~mask] 0 targets_onehot torch.nn.functional.one_hot( targets, num_classesnum_classes).permute(0, 3, 1, 2).float() denom (probs * mask.unsqueeze(1).float()).sum(dim(0, 2, 3)) numer (targets_onehot * (probs * mask.unsqueeze(1).float())).sum(dim(0, 2, 3)) dice (2 * numer self.smooth) / (denom targets_onehot.sum(dim(0, 2, 3)) self.smooth) return 1.0 - dice.mean()先将预测概率做成 one-hot 对齐目标mask剔除ignore_index对应的像素避免空洞标注干扰训练。两类分母分别统计预测区域面积与目标区域面积Dice 系数达到 1 时损失为 0。实际训练中我会把 DiceLoss 和交叉熵按 1:1 加权相加某些类别特别稀疏时再调高 Dice 的比重到 0.7。单纯用 Dice Loss 收敛后期容易出现轻微过拟合保留交叉熵能维持整体的分类边界。4.2 poly 学习率与优化器参数LiteSeg 在 ImageNet 预训练权重基础上微调一般用 SGD 比 Adam 最终精度更高尤其分割任务对收敛终点的细节比较敏感。推荐的优化器参数组合如下参数推荐值说明optimizerSGD配合动量项收敛更稳momentum0.9标准动量避免来回震荡weight_decay1e-4正则项过大导致欠拟合base_lr0.01batch size 16 时可尝试 0.02lr_schedulerpoly(1 - iter/total_iter)^0.9freeze_bn否微调时保持 BN 统计更新crop_size512显存小就降到 448 或 384深度分割任务的常用学习率策略是 poly 衰减而不是 StepLRpoly 策略在前半段保持较高学习率后半段平滑下降比阶梯式衰减更适合训练后期精细调整像素级输出。PyTorch 没有内置 poly 调度器需要自己写def poly_lr(base_lr, current_iter, max_iter, power0.9): return base_lr * (1.0 - current_iter / max_iter) ** power optimizer torch.optim.SGD( model.parameters(), lr0.01, momentum0.9, weight_decay1e-4) for epoch in range(epochs): for it, (images, masks) in enumerate(train_loader): lr poly_lr(0.01, len(train_loader) * epoch it, len(train_loader) * epochs) for g in optimizer.param_groups: g[lr] lr # 正常前反向传播power 取 0.9 是 DeepLab 系列一直沿用的经验值想加快收敛后期学习率下降速度可以调到 1.0。current_iter是全局迭代次数不是当前 epoch 内的循环变量这一点写错会导致学习率曲线整体偏移。如果你的训练周期很长也可以改用torch.optim.lr_scheduler.LambdaLR包装同样的函数这样 TensorBoard 能直接看到学习率变化。4.3 混合精度与显存不足的应对LiteSeg 本身参数不多但批量开到 16 到 32 后显存压力依然存在。PyTorch 的自动混合精度把前向计算切换到 FP16同时用 FP32 保存主权重能在不损失精度的前提下省约一半显存、提升约 40% 的训练吞吐。from torch.cuda.amp import GradScaler, autocast scaler GradScaler() for images, masks in train_loader: images, masks images.cuda(), masks.cuda() optimizer.zero_grad() with autocast(): outputs model(images) loss cross_entropy(outputs, masks) 0.5 * dice_loss(outputs, masks) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()autocast()只对包含卷积、矩阵乘法的算子切换 FP16BN 和损失函数仍然保持 FP32 精度避免数值溢出。GradScaler在反向传播前把 loss 放大若干倍梯度下溢为 0 时自动调整缩放因子。新版 PyTorch 已建议把torch.cuda.amp.autocast换成torch.amp.autocast(device_typecuda)但老代码用前者目前仍能运行升级后留意弃用警告即可。如果开了混合精度依然 Out of Memory最先做的是把crop_size从 512 降到 448分辨率下降 20% 左右显存占用按平方下降。其次把 batch size 降半并用梯度累积补打 batch sizeaccum_steps 2 for it, (images, masks) in enumerate(train_loader): loss loss / accum_steps loss.backward() if (it 1) % accum_steps 0: optimizer.step() optimizer.zero_grad()梯度累积让每个 step 等效于两倍 batch size 的梯度但 BN 的统计仍然基于单张图像batch size 降到 4 以下时 BN 均值方差抖动变大这时需要考虑换用 GroupNorm。最后一个手段是关闭主干网络的辅助深度监督输出辅助损失对主线精度提升有限但它会额外占几十 MB 显存和一份反向传播时间。5. 验证与部署加速测 FPS 的正确方式与量化落地5.1 用 CUDA Event 正确测量实时帧率训练完模型进入实时性验证环节。很多人在 PyTorch 里用time.time()循环测速得到的帧率极不稳定原因是 GPU 计算是异步的CPU 上的time.time()经常提前结束没有真正等到 GPU kernel 执行完成。正确的测速方式是用 CUDA Eventimport torch from torch.cuda import Event model.eval() x torch.randn(1, 3, 512, 512).cuda() for _ in range(10): with torch.no_grad(): _ model(x) torch.cuda.synchronize() start Event(enable_timingTrue) end Event(enable_timingTrue) start.record() with torch.no_grad(): for _ in range(100): _ model(x) end.record() torch.cuda.synchronize() avg_ms start.elapsed_time(end) / 100 print(f平均耗时: {avg_ms:.2f} ms, FPS: {1000 / avg_ms:.1f})前 10 次推理是 warm-up用于触发 CUDA kernel 和 BN 统计的初始化把真正测速时的额外开销排除掉。torch.cuda.synchronize()强制 CPU 等待 GPU 全部完成start.elapsed_time(end)返回毫秒数。这个方法测出来的是单张串行推理延迟如果部署时用生产级服务做批量并发FPS 还会更高。5.2 导出 ONNX 与静态量化的一条可行路径确认精度和速度达到预期后下一步通常是把模型导成 ONNX方便切换到 ONNX Runtime 或 TensorRT 推理。PyTorch 导出 ONNX 的步骤如下model.eval() dummy torch.randn(1, 3, 512, 512).cuda() torch.onnx.export( model, dummy, liteseg.onnx, input_names[input], output_names[output], opset_version12, dynamic_axes{input: {0: batch}, output: {0: batch}} )opset_version12是为了兼容性新版 ONNX Runtime 可以用 14 或 17。dynamic_axes把 batch 维度标记为动态部署时就能一次推理多张图。导出后用 ONNX Runtime 跑一遍验证输出确认和 PyTorch 的预测结果一致通常允许 1e-4 级别的浮点误差。如果产品对显存和带宽有硬性要求可以继续做 INT8 静态量化。这里给出一个可复现的量化技巧用校准集统计激活值范围而不是直接对权重做均匀量化。PyTorch 中通过torch.quantization.prepare插入观察算子在少量验证图片上跑一遍前向完成校准再convert得到量化模型。注意后端要用fbgemm在 GPU 上训练、CPU 上部署的场景下量化模型相对 FP32 通常能获得 2 到 4 倍加速。提示INT8 量化后必须在完整验证集上重新计算 mIoU与 FP32 模型对比。量化带来的精度回退通常应控制在 1 个百分点以内超出时需要尝试逐层敏感度分析把对量化最敏感的卷积层保留为 FP16。另一处是 BN 层在量化前必须折叠进卷积torch.quantization的 prepare 阶段会自动处理但你从 ONNX 转到 TensorRT 时要用带 BN 折叠的导出脚本否则推理结果和训练时可能对不上。本文还有配套的精品资源点击获取