LiteSeg轻量语义分割网络：PyTorch实现与部署实践

简介：LiteSeg实时轻量级语义分割算法的PyTorch实现，面向需要在边缘设备、低功耗硬件上完成实时推理的算法工程师与研究者，适用于自动驾驶、无人机监控、医疗影像分析等像素级分类场景。压缩包共39个文件，以21个Python源文件为主，覆盖模型定义、训练、评估与ONNX转换全流程；另有样例图片、配置文件和预训练权重，方便直接运行验证，整包约21.09MB。目前已有1358人下载学习。资源内置MobileNet、ShuffleNet、DarkNet等多种backbone变体，以及深度可分离卷积、ASPP、边缘增强等关键模块；同时提供数据增强、focal loss、在线难例挖掘、IOU评估、标签生成等工程化脚本，可直接替换数据集进行训练调优，也适合二次开发部署。目录结构清晰，代码注释规范，对入门轻量分割和工程落地都很有参考价值。

1. 实时语义分割的算力困境与 LiteSeg 的切入点

在嵌入式设备、移动端 SoC 上做逐像素分类，最常见的矛盾是：精度达标的模型跑不动，能跑起来的模型分割边界又糊成一团。DeeplabV3 这类经典架构在 1080Ti 上可以做到实时，但换到 Jeston Nano 或者手机 NPU 上，单帧延迟立刻涨到几百毫秒，根本无法接入视频流。LiteSeg 正是冲着这个场景去的——在 PyTorch 框架下用轻量化主干替换标准 ResNet，配合空洞空间金字塔池化保留多尺度感受野，在参数量和 mIoU 之间找平衡点。它的 v2 版本把主干可选范围扩展到 MobileNet、ShuffleNet、DarkNet 三系，配套训练、评估、ONNX 导出脚本齐全，适合两类人：一类是刚入门语义分割、想研究轻量网络怎么设计的研究生；另一类是已经在做边缘端部署、需要快速评估不同主干精度与延迟的工程师。

2. LiteSeg 网络架构与轻量化设计选型

2.1 编码器解码器框架与轻量主干的组合逻辑

LiteSeg-v2 源码仓库里的结构很清晰：liteseg.py是整体的前向逻辑，backbone_networks目录下分别给出了liteseg_mobilenet.py、liteseg_shufflenet.py和liteseg_darknet.py，三个文件各自的命名暗示了主干网络的三种选择。这种设计思路是标准的两段式语义分割模型：编码器负责从输入图像提取高维语义特征，解码器负责将低分辨率特征图恢复成原分辨率预测。解码端在这里并不是论文里那种复杂的逐步上采样，而是依赖aspp.py提供的多尺度空洞卷积特征聚合——整体结构参考了 DeepLab 系列的思路，但把主干替换成了推理开销更小的分类网络。

选择轻量主干的时候，首先要确认的是替换的成本边界。MobileNetV2 的倒残差结构在 ImageNet 上大约 70% 出头的 top-1 精度，比 ResNet50 低 5 个点左右，但 FLOPs 不到 ResNet50 的六分之一。这个精度的下降在分割任务上体现为特征图语义信息密度的降低，最终影响 mIoU；但换来的是视频流场景能跑满帧率。具体到 LiteSeg 的代码，liteseg_shufflenet.py里的 ShuffleNetV2 利用 channel shuffle 增强跨通道信息交互，它的优势在于实际推理速度，因为逐通道卷积在底层硬件上更容易被优化。DarkNet 分支则偏向精度，参数量介于前两者之间。

2.2 ASPP 空洞空间金字塔池化的参数设计

aspp.py是全项目的核心组件，它负责在多个空洞率下并行采样特征。常见做法是在 ASPP 模块里设置四个分支：一个 1x1 卷积做降维、三个 3x3 空洞卷积分别使用 6、12、18 的空洞率，再加一个全局平均池化分支。源文件aspp.py中这几个空洞率的取值直接决定了感受野覆盖范围。空洞率 6 的分支侧重局部细节，12 覆盖中等范围的上下文，18 负责捕捉大目标整体结构，最后池化分支提供整图级别的统计信息。

这里有一个在资源包在训练和推理时需要特别注意的参数设置：输入分辨率变化时，ASPP 的感受野是等比放大的。inputs尺寸如果从 512x512 改到 768x768，空洞率 18 对应的采样范围增加，对大一倍的目标依然有覆盖，但小目标的分割可能因为背景信息过多而出现误分类。跑train.py之前如果自己改了crop_size，建议同步检查 aspp dilation 是否需要缩放。

2.3 深度可分离卷积的实现考量

separableconv.py提供了逐通道卷积加逐点卷积的封装，这是 MobileNetV2 的基础构建单元。一个标准 3x3 卷积在输入输出通道均为 C 时参数是 9 * C * C，换成深度可分离卷积后，逐通道卷积参数 9 * C，逐点卷积参数 C * C，总参数量从 9C² 降到 C² + 9C，当 C 是 256 时计算量约为原来的七分之一。源码里需要留意的是逐通道卷积的参数组设置，PyTorch 中实现时groups必须等于输入通道数：

# separableconv.py 的核心结构示意 import torch.nn as nn class SeparableConv2d(nn.Module): def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1): super().__init__() # depthwise 卷积，groups=in_channels 意味着每个输入通道单独卷 self.depthwise = nn.Conv2d( in_channels, in_channels, kernel_size=kernel_size, stride=stride, padding=padding, groups=in_channels, bias=False ) # pointwise 卷积，只做 1x1 通道融合 self.pointwise = nn.Conv2d( in_channels, out_channels, kernel_size=1, bias=True ) def forward(self, x): return self.pointwise(self.depthwise(x))

每组卷积groups=in_channels意味着没有跨通道的信息融合，这使得模型前几层的语义抽象能力变弱。在训练 LiteSeg 时如果发现收敛速度明显变慢，优先检查separableconv.py里 depthwise 卷积之后是否接了 BatchNorm 和激活函数，这比标准卷积更依赖 BN 来稳定分布。

2.4 三种主干的取舍建议

仓库liteseg_darknet.py与liteseg_mobilenet.py在精度上的差距大约在 2 到 4 个 mIoU，但 DarkNet 系列的推理耗时会高出 30% 以上。以下是对比维度：

区分这三种结构有一个快速判断标准：看backbone_networks中对应文件里Conv2d的groups参数，如果groups等于输入通道数就是深度可分离变体，等于输入通道数一半则通常是 ShuffleNet 风格的分组卷积。这三条路并不冲突，实际项目里可以先用 MobileNet 跑通训练流程，再用 DarkNet 微调验证上限。

3. PyTorch 训练管线：数据增强、损失函数与训练配置

3.1 augmentation 与数据集加载的细节实现

常规分割训练里的增强策略与检测任务不同，不能随便做随机的方形裁剪，因为标签图必须和输入做完全相同的空间变换。dataloaders/augmentation.py中需要同时处理 image 和 mask 两个数组，最常用的组合是随机水平翻转、尺度抖动和颜色抖动。我在实际跑训练时会固定使用以下配置：随机缩放比例在 0.5 到 2.0 之间，缩放后再随机裁剪到固定输入尺寸，最后以 0.5 概率水平翻转。这样做的好处是模型对目标尺度变化不敏感，在道路场景中物体大小差异非常大，纯固定分辨率训练出来的模型换到新环境往往要重新调。

彩色图像的均值方差统计对收敛速度有直接影响。PyTorch 的transforms.Normalize接收三个通道的 mean 和 std，LiteSeg 的utils.py里如果沿用 ImageNet 的均值[0.485, 0.456, 0.406]，加载自己采集的数据集时也需要保持。这里有一个很多人忽略的坑：使用 ImageNet 预训练权重的话，输入归一化标准必须和预训练时保持一致，否则浅层特征分布错位，前几个 epoch 的 loss 会跳得厉害且很难恢复。

# 数据集目录结构建议 datasets/ ├── train.txt # 训练样本相对路径列表 ├── val.txt # 验证样本相对路径列表 ├── JPEGImages/ # 原图 └── SegmentationClass/ # 伪彩色标签 或 灰度标签

如果训练自己的数据，parse_classes.py负责把类别名称映射到固定索引，注意类别顺序一旦确定就不要改动，否则评估脚本iou_eval.py计算出的类别 AP 全部错位。

3.2 损失函数选择与 OHEM 的适用边界

项目里loss.py的实现决定训练收敛走向。语义分割最常见的是交叉熵损失，但城市道路场景里类别极度不均衡，比如树木和天空占比很大、摩托车和交通标志占比很少。直接用交叉熵会导致模型把所有像素预测为占比最大的类别也能拿到很低的 loss。此时有两种主流改进方案：一是用focal loss降低易分类样本的权重，二是用在线难样本挖掘 OHEM，两者处理的侧重不同。OHEM 的做法是每个 batch 计算完所有像素的损失之后，按损失值从大到小排序，只取前 N 个像素做反向传播。这个 N 通常设为 batch 内总像素数的一半，也就是 50% 的 top-k 比例。在源代码中可以通过判断 loss 的topk部分完成，但要注意 OHEM 有一个副作用：如果某个类别正好在难样本中占绝大比例，模型可能被带偏，变得更关注那一个类别。

训练时如果发现验证集上某个小类别完全预测不出来，先关闭 OHEM 换成带类别权重的交叉熵，权重的设置公式一般取出现频率的倒数再开根号，避免权重值压得过于极端。

3.3 train.py 训练流程与核心参数调整

仓库里的train.py是入口脚本，配置信息分离在config/training.yaml中，通过environment.yml可以一键创建可复现的 conda 环境。重点参数集中在training.yaml的batch_size、base_lr、crop_size、epochs和backbone这几项。按照经验，输入crop_size=512时 batch size 在 8 到 16 之间可以占满一块 11GB 显存的显卡。学习率策略常用 poly 衰减，即初始学习率乘以(1 - iter / total_iters) ** 0.9的幂指数，在末尾阶段学习率降到接近 0，让模型充分收敛。

以下是我修改过的一个典型配置片段，可以对照自己的显存调整：

# config/training.yaml 关键字段 model: backbone: mobilenet # 可选 mobilenet / shufflenet / darknet aspp_rates: [6, 12, 18] output_stride: 16 # 编码器下采样倍数，8 更慢但边界更好 train: crop_size: 512 batch_size: 12 base_lr: 0.01 # 配合 poly 策略 power: 0.9 epochs: 120 warmup_epochs: 3 # 前几个 epoch 线性升 lr，防止初期振荡

output_stride指的是输入到最终特征图的空间下采样倍数。设为 16 时特征图是输入的 1/16，对应空洞卷积的空洞率可以进行微调；设为 8 时精度提升但显存用量接近翻倍。在足够显存的显卡上跑语义分割时，如果目标是追求更高的 mIoU，output_stride=8的性价比很高；如果目标是出 demo、跑视频流，16 更现实。

3.4 训练过程的监控指标

iou_eval.py提供了逐类别的 IoU 计算，训练中间产生的 checkpoint 通过utils/card.py生成可视化对比图。checkpoints/目录下有现成的 LiteSeg-mobilenet-card_epoch-99.pth 权重文件，这个权重可以直接加载跑demo.py看效果。评估自己的模型时一定要同时看mIoU和class IoU的每个类别单独数值，mIoU 被几个大类拉高的情况在长尾分布的数据集上极其常见，只报一个总分会掩盖部分类别完全不可用的事实。

4. 导出部署流程：从 PyTorch 权重到端侧推理

4.1 ONNX 导出的输入输出约定

convert_to_onnx.py把训练好的权重转换为中间表示，这是进入推理框架前的标准一步。PyTorch 模型需要锁定输入尺寸和 batch 维度后再导出，因为 ONNX 文件里的 shape 信息对后续 TensorRT 或 OpenVINO 的优化起决定性作用。以下是导出时的常见参数设定：

python convert_to_onnx.py \ --checkpoint checkpoints/LiteSeg-mobilenet-card_epoch-99.pth \ --backbone mobilenet \ --height 512 \ --width 512 \ --opset 11

opset版本决定转换器允许使用的算子集合。PyTorch 的torch.onnx.export里动态轴dynamic_axes虽然能导出可变分辨率模型，但 TensorRT 会在运行期做形状推断，动态 batch 的优化效果明显弱于固定 shape。在视频流应用里输入的分辨率通常是稳定的，固定 h/w/batch 是更优做法，前提是模型的crop_size尽量和部署尺寸一致，避免训练与推理分辨率不一致导致的精度掉点。

4.2 导出的输出 opset 版本与算子兼容

转换完成后，输出的 logits 形状是(1, num_classes, H, W)。常见推理框架的输入要求是固定的NCHW格式，颜色顺序必须为 RGB。导出的过程如果遇到不支持的算子，处理优先级是：先换更高的opset版本，PyTorch 官方对算子的支持力度随着 opset 新版本持续增强；如果仍然失败，把对应模块在模型代码里改写成算子友好的等价形式。

4.3 TensorRT 加速与 INT8 量化

转换得到 ONNX 后，部署到英伟达的 GPU 平台可以用 TensorRT 做落地推理，这也是资源包在 eval 环节最实用的加速手段。TensorRT 在 FP16 精度下通常可以获得 2 到 3 倍的加速比，显存占用减小约一半。如果对精度损失敏感，可以只对depthwise卷积层做 INT8 量化，其余层保持 FP16。量化校准数据集一般取验证集里随机采样的 500 张图，在校准时需要统计每个激活值的动态范围。INT8 推理对边界区域的输出影响最明显，量化之后物体边缘往往会出现锯齿，这是通道级量化无法避免的信息损失。

# TensorRT Python API 推理核心流程 engine = builder.build_engine(network, config) context = engine.create_execution_context() context.set_binding_shape(0, (1, 3, 512, 512)) # 输入绑定 output = np.empty((1, num_classes, 512, 512), dtype=np.float32) context.execute_v2([input_ptr, output_ptr])

TensorRT 的构建阶段会做算子融合和层间内存复用，但要注意模型里如果有torch.argmax这类非网络层操作不要直接放进导出范围，后处理留在宿主代码里做。实际部署中 v2 版本直接运行好。

5. 验证技巧与训练排错经验

5.1 边界质量的量化验证

语义分割评估如果只看 mIoU，很容易忽略边界模糊的问题。mIoU 是通过区域重叠计算的指标，对边界处 1 到 2 个像素的偏移不敏感，但对视觉观感来说恰好是这些边缘位置的锯齿感和漏检最明显。此时可以用边界 F1 分数来补充验证：把真实标签和模型预测都做 Canny 边缘提取，再计算边缘像素的 Precision、Recall 和 F1。实现这个验证逻辑不复杂：先对预测图做 one-hot 转换，再逐类别计算边界图，最后以几个像素宽度的容差做匹配。建议在评估脚本里单独加一组边界 IoU 统计，观察它在不同主干下的变化趋势。以我验证 Mo-bileNet 主干和 DarkNet 主干的实践反馈，前者的 mIoU 可能只掉两个点，但边界 F1 会显著下降，这个差距直接反映在人的视觉印象上。

5.2 训练崩溃的常见原因排查

训练在某个 epoch 突然出现 NaN 是常见问题。第一条定位路径是查看training.yaml里的base_lr，很可能是学习率过大加上poly衰减末期的数值振荡造成的。可以把初始学习率降到原来的十分之一重新训练，排查是否还出现 NaN。如果是权重初始化或者 BatchNorm 的 momentum 问题，加载checkpoints/目录下的预训练权重继续训练能规避这类风险。

显存不足时优先调整的是crop_size而不是batch_size。crop_size减半会带来 4 倍的显存节约，同时面积减半意味着每个 batch 内的样本信息量变小，可以同时把batch_size适当调大填平训练效率。在数据加载部分有参数量较大预算时可尝试激活torch.backends.cudnn.benchmark = True加速卷积搜索。训练完成后做推理加速时，优先把输入图像通过cv2.resize等比缩放到模型输入规格，而不是直接裁剪，保留宽高比能明显减少小目标丢失。

最后在推理侧检查模型速度时，确认 batch size 为 1、并且关了autograd梯度记录，用torch.no_grad()包裹推理代码，这是所有人最常忘记的一步：开着梯度推理时显存占用高且计算会额外记录反向图，实测耗时差距可能在 30% 以上。

本文还有配套的精品资源，点击获取