简介:这是一个基于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 python=3.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_index=255处理,后面训练会用到。
制作自定义数据集时,常见做法是用 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, transform=None): 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, dtype=np.int64)) return image, mask这里的核心点是:掩码不能走torchvision.transforms.ToTensor(),因为ToTensor()会把像素值归一化到 0 到 1,导致标签变成浮点数,而CrossEntropyLoss要求标签是整数张量。直接np.array(mask, dtype=np.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, stride=1): return nn.Sequential( nn.Conv2d(inp, oup, 3, stride, 1, bias=False), nn.BatchNorm2d(oup), nn.ReLU6(inplace=True) ) 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, stride=1)) layers.extend([ nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False), nn.BatchNorm2d(hidden_dim), nn.ReLU6(inplace=True), nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False), 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)groups=hidden_dim就是深度可分离卷积里的 depthwise 部分,每个通道单独做卷积,参数量从hidden_dim * 9降到hidden_dim * 1。expand_ratio一般取 6,也就是中间维度是输入通道的 6 倍。stride为 2 时下采样并且不启用残差连接,因为特征图尺寸变了不能直接相加。
组装 MobileNetV2 时需要打一个表,按输入分辨率递减的顺序堆叠各阶段。通常在最后一个 block 输出的特征图尺寸是输入图像的 1/16 或 1/32,对应output_stride=16或 32。LiteSeg 希望保留稍高的分辨率来挽救小目标,所以常见做法是不在最后一个阶段做 stride=2 的下采样,改用它后面的空洞卷积保持感受野。实际实现时,可以把原本 stride=2 的 block 改成 stride=1,再把同一阶段里所有 3x3 depthwise 卷积设置 dilation=2,这样特征图停在 1/16 分辨率。
3.2 LR-ASPP:LiteSeg 的多尺度语义聚合
ASPP 的设计思路是在同一个特征图上用不同空洞率的并行空洞卷积采样,空洞率越大感受野越大,从而同时捕获小目标和大物体。标准 DeepLabV3 里的 ASPP 使用空洞率 6、12、18 的 3x3 卷积,外加一个全局平均池化分支,在 Cityscapes 上效果不错但计算量偏大。LiteSeg 把它改成 LR-ASPP,核心改动是去掉 256 通道的大卷积,每个空洞分支只输出 128 通道,并且把 3x3 卷积换成深度可分离卷积。
class LRASPP(nn.Module): def __init__(self, in_channels, out_channels=128, 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, padding=rate, dilation=rate, bias=False), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True) )) self.image_pool = nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Conv2d(in_channels, out_channels, 1, bias=False), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True) ) self.fuse = nn.Sequential( nn.Conv2d(out_channels * (len(atrous_rates) + 1), out_channels, 1, bias=False), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True) ) 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, size=x.shape[2:], mode="bilinear", align_corners=False) branch_outs.append(pool_out) return self.fuse(torch.cat(branch_outs, dim=1))dilation=rate决定了卷积核采样的间距,rate 为 6 时 3x3 卷积的实际感受野等效 13x13,rate 为 18 时等效 37x37。使用padding=rate保证输出特征图尺寸不变,方便后续 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, bias=False), nn.BatchNorm2d(64), nn.ReLU(inplace=True), 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, size=low_level_feat.shape[2:], mode="bilinear", align_corners=False) x = torch.cat([high_level_feat, low_level_feat], dim=1) return self.decoder(x)这里的low_level_feat来自 MobileNetV2 的浅层,分辨率是输入的 1/4,high_level_feat来自深层,分辨率是 1/16。interpolate把深层的 1/16 特征上采样到 1/4,然后与浅层特征拼接。decoder里第一个 3x3 卷积的输入通道是128 + low_level_feat.channels,要根据实际 backbone 输出调整。如果浅层特征取 24 通道,拼接后就是 152 通道。
如果你不想手写 backbone,torchvision.models里有官方 MobileNetV2,可以加载在 ImageNet 上预训练的权重做迁移学习。LiteSeg 训练时微调编码器权重能明显加速收敛,比从头训练节省一半以上时间。
3.4 验证轻量级:模型参数和计算量的度量
构建完网络后先量化指标,确认它真的符合实时轻量级的定位。用 thop 库统计参数量和 FLOPs:
pip install thopfrom thop import profile, clever_format model = LiteSeg(backbone, aspp, num_classes=21).cuda() model.eval() flops, params = profile(model, inputs=(torch.randn(1, 3, 512, 512).cuda(),)) flops, params = clever_format([flops, params], "%.3f") print(f"FLOPs: {flops}, Params: {params}")FLOPs 表示模型处理一张 512x512 图像所需的浮点运算次数,Params 是模型存储参数量。LiteSeg 在 MobileNetV2 backbone、输出 stride 16 的情况下,参数量大约在 4M 到 8M 之间,比 DeepLabV3 的 40M 到 60M 小一个数量级。统计时注意 FLOPs 通常不包含 BN 和激活函数的计算,不同工具计算口径有差异,所以我一般只拿同一工具对比不同模型的相对大小,而不是跟论文里数字绝对比较。
4. 训练与调参:损失函数、学习率与显存策略
4.1 损失函数与类别不平衡
语义分割最常用的损失是逐像素交叉熵,但直接用它处理 Cityscapes 这类类别分布极不均匀的数据集时,模型会偏向出现频率高的类别。LiteSeg 的常见做法是交叉熵加 Dice Loss,Dice Loss 直接优化预测区域和真实标注区域的重合度,对像素数量少的类别更敏感。
class DiceLoss(nn.Module): def __init__(self, smooth=1.0, ignore_index=255): super().__init__() self.smooth = smooth self.ignore_index = ignore_index def forward(self, logits, targets): num_classes = logits.shape[1] probs = torch.softmax(logits, dim=1) mask = targets != self.ignore_index targets = targets.clone() targets[~mask] = 0 targets_onehot = torch.nn.functional.one_hot( targets, num_classes=num_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 最终精度更高,尤其分割任务对收敛终点的细节比较敏感。推荐的优化器参数组合如下:
| 参数 | 推荐值 | 说明 |
|---|---|---|
| optimizer | SGD | 配合动量项收敛更稳 |
| momentum | 0.9 | 标准动量,避免来回震荡 |
| weight_decay | 1e-4 | 正则项,过大导致欠拟合 |
| base_lr | 0.01 | batch size 16 时可尝试 0.02 |
| lr_scheduler | poly | (1 - iter/total_iter)^0.9 |
| freeze_bn | 否 | 微调时保持 BN 统计更新 |
| crop_size | 512 | 显存小就降到 448 或 384 |
深度分割任务的常用学习率策略是 poly 衰减而不是 StepLR,poly 策略在前半段保持较高学习率,后半段平滑下降,比阶梯式衰减更适合训练后期精细调整像素级输出。PyTorch 没有内置 poly 调度器,需要自己写:
def poly_lr(base_lr, current_iter, max_iter, power=0.9): return base_lr * (1.0 - current_iter / max_iter) ** power optimizer = torch.optim.SGD( model.parameters(), lr=0.01, momentum=0.9, weight_decay=1e-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()只对包含卷积、矩阵乘法的算子切换 FP16,BN 和损失函数仍然保持 FP32 精度,避免数值溢出。GradScaler在反向传播前把 loss 放大若干倍,梯度下溢为 0 时自动调整缩放因子。新版 PyTorch 已建议把torch.cuda.amp.autocast换成torch.amp.autocast(device_type='cuda'),但老代码用前者目前仍能运行,升级后留意弃用警告即可。
如果开了混合精度依然 Out of Memory,最先做的是把crop_size从 512 降到 448,分辨率下降 20% 左右,显存占用按平方下降。其次把 batch size 降半,并用梯度累积补打 batch size:
accum_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 Event:
import 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_timing=True) end = Event(enable_timing=True) 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_version=12, dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}} )opset_version=12是为了兼容性,新版 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 折叠的导出脚本,否则推理结果和训练时可能对不上。
本文还有配套的精品资源,点击获取