做语义分割项目的时候很多人的第一反应是抱一个 ResNet 或者 VGG 系列的编码器再套上 U-Net 解码器。跑通很容易但一上实际部署就发现参数太多、推理太慢、显存吃紧。我自己在做一套基于 PyTorch 的图像分割全套项目时刻意把编码器换成了 MobileNetV1配合 U-Net 解码器最终在保持不错分割精度的前提下把模型参数量压缩到了一个非常舒服的量级。这篇文章就把这套 MobileNetV1-U-Net 语义分割项目的完整思路、关键代码、训练配置和踩坑记录全部拆开讲清楚适合已经跑通过基础分割网络、想进一步做轻量化落地的开发者参考。这套项目不是我临时拼凑的 Demo而是完整包含了数据预处理、网络定义、训练脚本、评估指标计算和模型导出。下面我从选型原因开始一层层把网络结构、代码实现、训练细节讲透。如果你也在纠结轻量骨干网络到底能不能撑起分割任务这篇文章应该能给你一个比较实在的答案。1. 轻量编码器选型MobileNetV1 在语义分割中的底气来源1.1 语义分割项目最终要什么语义分割本质上是对图像的每个像素做分类。它和图像分类最大的区别在于分类只需要一个全局语义向量而分割必须同时保留空间位置信息。这就是为什么纯分类网络不能直接用于分割——最后一层特征图已经被压到 1/32 分辨率空间信息大量丢失需要解码器逐步恢复。做轻量化分割项目时我给自己定了三个目标参数量控制在 10M 以下单张 512x512 图像的前向推理延迟控制在可接受范围同时 MIOU 不能比 VGG 骨干的版本差太多。从这三个目标出发骨干网络的选择就非常关键了。ResNet18 虽然参数也不算多但标准残差结构在低算力设备上的并行效率不如深度可分离卷积MobileNetV3 精度更好但结构里引入了 SE 模块和 h-swish在部署时对算子支持的要求更高。反倒是 MobileNetV1结构极度简洁只有深度卷积和逐点卷积两种算子几乎所有推理引擎都能无缝支持。1.2 从 VGG 到 MobileNet 的替换动机如果你用过 VGG16 做 U-Net 编码器应该对那个 134M 的参数量记忆犹新。VGG16 前面几层都是 64、128 通道的 3x3 卷积虽然 FLOPs 不算极端但参数量太大训练完的模型文件动辄 500MB。MobileNetV1 把标准卷积拆成了两步先用深度卷积在单个通道内做 3x3 卷积然后用 1x1 逐点卷积做通道间的线性组合。这样一来参数量从标准卷积Dk * Dk * M * N变成深度可分离卷积Dk * Dk * M M * N其中 Dk 是卷积核边长M 是输入通道数N 是输出通道数。当 Dk3、M 和 N 都是 256 时标准卷积参数量是 589,824深度可分离卷积只有 2,304 65,536 67,840直接少了将近 8 倍。这个压缩比放到整个编码器上效果非常明显。我用 MobileNetV1 替换 VGG16 做编码器后整体模型参数量从大约 134M 降到了约 7.4M。训练时显存占用大幅下降单卡 11GB 显存可以轻松塞下 batch size 8 的 512x512 输入。这个数字我在实测中验证过不同 PyTorch 版本下统计方式会略有差异但基本都在这个量级。1.3 MobileNetV1 的固有缺陷与绕过方式MobileNetV1 不是没有缺点。它没有残差连接深层的梯度传导完全靠普通卷积堆叠在训练初期容易出现梯度不稳定。尤其是用大批量训练时如果 BN 参数不处理好loss 可能直接崩掉。我的处理方式是迁移 ImageNet 预训练权重作为初始化训练阶段冻结编码器前两个 stage 的 BatchNorm 参数。这样既保留了浅层的通用边缘和纹理特征又避免了小 batch 下 BN 统计量剧烈波动。这个方案在后面的训练章节会详细展开。2. MobileNetV1 编码器拆解特征图到底下采样了多少次2.1 编码器的特征图输出节点MobileNetV1 的分类网络输入是 224x224连续经过 13 个深度可分离卷积模块每次 stride2 的模块将分辨率减半。整体下采样倍率是 32对应输出 7x7 特征图。我把编码器按下采样节点拆成了 4 个 stage每个 stage 的输出都会保留下来作为给解码器的跳跃连接stage1stride2 一次输出 112x112通道 64stage2stride2 一次输出 56x56通道 128stage3stride2 一次输出 28x28通道 256stage4stride2 一次输出 14x14通道 512stage5stride2 一次输出 7x7通道 1024这里有个值得说的选择很多语义分割文章会把最后一个 stage 的 stride 改成 1配合空洞卷积把下采样倍率维持在 16。这样做理论上能保留更密集的特征但会让编码器输出分辨率变大解码器随之变重推理变慢。我实测下来对于室内场景、遥感地物这类边界相对规整的数据集stride32 的编码器配合完整解码器已经够用省下的算力可以留给输入分辨率。所以我保留经典的 32 倍下采样结构没有做空洞卷积改动代码也更简洁。如果你做的是细胞膜边界、车道线这类极细边缘的分割任务可以把 stage4 和 stage5 的 stride 改成 1 再补空洞卷积MIOU 在细边界的 Recall 上会有几个点的提升。2.2 为什么 MobileNetV1 浅层特征对分割很重要U-Net 最大的贡献就是跳跃连接。浅层特征图分辨率高、通道数少保存了图像边缘、颜色、纹理这些局部信息深层特征图分辨率低但语义强能分辨这个东西是桌子还是椅子。在 MobileNetV1-U-Net 结构里stage1 的 112x112 特征图至关重要。它保留了原始输入一半分辨率的空间细节解码器最后两级恢复分辨率时主要就是靠它来精修边界。我建议你在看一张特征图可视化的时候对比一下 stage1 和 stage5 的激活图。stage5 的倒数第二个通道通常能看出完整的物体轮廓而 stage1 的某个通道可能在响应一个非常细的纹理方向。这个对比会让你直观理解跳跃连接存在的意义。3. U-Net 解码器设计从 1/32 分辨率找回边缘细节3.1 上采样方式选型双线性插值还是转置卷积解码器要做的事是把 7x7 的特征图一步步恢复到输入分辨率。上采样无非就两条路线转置卷积和插值卷积。我最终选了双线性插值 3x3 卷积的组合而不是转置卷积。原因有三第一转置卷积有比较明显的棋盘效应风险。当 kernel_size 不能被 stride 整除时输出的某些位置会周期性重叠形成让人头疼的伪影。双线性插值不存在这个问题。第二转置卷积的参数量大。每层转置卷积都要学习权重放在轻量化项目里不划算。第三插值本身没有可学习参数上采样的分辨率恢复逻辑是固定的配合后面的 3x3 卷积学习特征融合效果更稳。实际操作中我用nn.Upsample(scale_factor2, modebilinear, align_cornersFalse)做 2 倍上采样然后接一个 3x3 卷积完成通道变换和特征融合。PyTorch 实测这个组合既稳定又快。3.2 跳跃连接怎么接才不会让解码器膨胀MobileNetV1 的编码器通道数到 stage5 已经是 1024。如果解码器每层都把跳跃连接的通道直接 concat 进来第一层解码器的输入通道会变成 1024 512 1536再卷一层输出 512参数量瞬间爆炸。我在解码器里做了通道压缩每个跳跃连接先过一个 1x1 卷积把编码器特征压缩到计划通道数的四分之一再做 concat。整个解码器的通道走势是512 - 256 - 128 - 64 - 32具体的 concat 逻辑是这样假设当前解码层上采样后的特征图是 F通道数 C对应编码器的跳跃连接特征图是 E通道数 C*4。我先把 E 用 1x1 卷积压到 C然后 concat 成 2C再过 3x3 卷积输出 C。这样做参数量的增加非常有限每一层解码器几乎都保持和上一层相同量级的计算量。最终整个模型的参数量才能稳定在 7.4M。3.3 最后一层输出头的细节解码器输出的是 32 通道的特征图最后接一个 1x1 卷积把通道映射到类别数 N。如果你的数据集有背景、地面、墙面、物体四个类别就是映射到 4。推理时对这个输出做 argmax 得到每个像素的类别索引再映射到调色板生成可视化结果。训练时则把输出直接送进交叉熵损失。这里有个小细节最后的 1x1 卷积不要加激活函数和 BatchNorm直接裸输出 logits。因为分割的损失函数要求的是未归一化的类别得分激活函数会破坏这个分布。4. PyTorch 实现详解从加载预训练权重到完整推理链路4.1 项目文件组织这套项目的代码结构比较清晰基本是每个环节一个独立模块config.py # 全局配置路径、类别、图片尺寸、超参数 models/ encoder.py # MobileNetV1 编码器 decoder.py # U-Net 解码器 seg_model.py # 编码器解码器整体网络 utils/ dataset.py # 数据读取与增强 metrics.py # MIOU / PA 计算 visualize.py # 预测结果可视化 utils.py # 权重加载、日志等工具 train.py # 训练主脚本 predict.py # 单图推理脚本 export.py # TorchScript 导出脚本拿到项目第一步肯定是改 config.py。别的可以先不动但DATA_ROOT路径必须指向你自己的数据集NUM_CLASSES必须和标签类别一致。很多同学跑我的项目报维度错误十有八九是这里没改。4.2 编码器核心代码MobileNetV1 的编码器实现核心在_build_stages。我用 PyTorch 的nn.Sequential按通道划分阶段同时在 forward 里把各阶段输出保存成 listimport torch import torch.nn as nn class MobileNetV1Encoder(nn.Module): def __init__(self, width_mult1.0, pretrainedTrue): super().__init__() def conv_bn(inp, oup, stride): return nn.Sequential( nn.Conv2d(inp, oup, 3, stride, 1, biasFalse), nn.BatchNorm2d(oup), nn.ReLU(inplaceTrue) ) def conv_dw(inp, oup, stride): return nn.Sequential( nn.Conv2d(inp, inp, 3, stride, 1, groupsinp, biasFalse), nn.BatchNorm2d(inp), nn.ReLU(inplaceTrue), nn.Conv2d(inp, oup, 1, 1, 0, biasFalse), nn.BatchNorm2d(oup), nn.ReLU(inplaceTrue) ) base_channels [32, 64, 128, 256, 512, 1024] if width_mult ! 1.0: base_channels [int(c * width_mult) for c in base_channels] self.stage1 nn.Sequential( conv_bn(3, base_channels[0], 2), conv_dw(base_channels[0], base_channels[1], 1) ) self.stage2 nn.Sequential( conv_dw(base_channels[1], base_channels[2], 2), conv_dw(base_channels[2], base_channels[2], 1) ) self.stage3 nn.Sequential( conv_dw(base_channels[2], base_channels[3], 2), conv_dw(base_channels[3], base_channels[3], 1) ) self.stage4 nn.Sequential( conv_dw(base_channels[3], base_channels[4], 2), conv_dw(base_channels[4], base_channels[4], 1) ) self.stage5 nn.Sequential( conv_dw(base_channels[4], base_channels[5], 2), conv_dw(base_channels[5], base_channels[5], 1) ) if pretrained: self._load_pretrained_weights() def _load_pretrained_weights(self): # 注意加载 ImageNet 预训练权重时必须剔除分类全连接层 pass def forward(self, x): c1 self.stage1(x) c2 self.stage2(c1) c3 self.stage3(c2) c4 self.stage4(c3) c5 self.stage5(c4) return [c1, c2, c3, c4, c5]加载预训练权重时有一个关键细节。MobileNetV1 的官方权重里最后有一个全连接分类层fc它的形状和语义分割任务不匹配必须剔除。加载时用strictFalse或者按 key 过滤def _load_pretrained_weights(self): import torchvision.models as models state_dict models.mobilenet_v1(weightsmodels.MobileNet_V1_Weights.IMAGENET1K_V1).state_dict() filtered {k: v for k, v in state_dict.items() if not k.startswith(fc)} self.load_state_dict(filtered, strictFalse)常见的坑是直接model.load_state_dict(state_dict)会报 unexpected key 的错误因为编码器的stage1.stage1这种命名方式和官方 mobilenet 的features命名对不上。我的做法是把官方权重里的 key 名做一个映射把features.0对应到stage1.0features.12对应到stage5.1。你也可以直接加载后再手动拆权重总之保证 key 完全匹配。4.3 解码器核心代码解码器的UpBlock是我最常用的模块逻辑很简单双线性上采样 2 倍然后和压缩后的跳跃连接 concat再过 3x3 卷积import torch import torch.nn as nn class UpBlock(nn.Module): def __init__(self, in_ch, skip_ch, out_ch): super().__init__() self.skip_conv nn.Conv2d(skip_ch, out_ch, 1, biasFalse) self.conv nn.Sequential( nn.Conv2d(in_ch out_ch, out_ch, 3, padding1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue) ) def forward(self, x, skip): x nn.functional.interpolate(x, scale_factor2, modebilinear, align_cornersFalse) skip self.skip_conv(skip) x torch.cat([x, skip], dim1) return self.conv(x) class MobileNetV1UNet(nn.Module): def __init__(self, num_classes4): super().__init__() self.encoder MobileNetV1Encoder(width_mult1.0, pretrainedTrue) # 解码器这里和编码器特征通道严格对应 self.up1 UpBlock(1024, 512, 512) self.up2 UpBlock(512, 256, 256) self.up3 UpBlock(256, 128, 128) self.up4 UpBlock(128, 64, 64) self.up5 UpBlock(64, 32, 32) self.head nn.Conv2d(32, num_classes, 1) def forward(self, x): c1, c2, c3, c4, c5 self.encoder(x) x self.up1(c5, c4) x self.up2(x, c3) x self.up3(x, c2) x self.up4(x, c1) x self.up5(x, x) # 最后一次无额外跳跃或者接一个池化特征 return self.head(x)注意我在 forward 的最后一层self.up5(x, x)处做了一个偷懒的处理最后一次上采样没有新的编码器特征可拼接就把自身 concat 了一次。你完全可以根据自己的数据集去掉这一层直接上采样到原分辨率后接 head。这个设计并不会带来额外的语义信息只是保持了代码结构的统一。4.4 预处理和后处理链路推理时的预处理细节不能省。MobileNetV1 是在 ImageNet 上预训练的输入必须按照 ImageNet 的均值方差做 normalize。很多同学后端接了一个 resize 就喂给网络结果分割边界全是噪点大概率就是这里出了问题def preprocess(image: np.ndarray) - torch.Tensor: # 输入是 HWC 的 BGR/ RGB 图像 image cv2.resize(image, (512, 512)) image image.astype(np.float32) / 255.0 mean np.array([0.485, 0.456, 0.406], dtypenp.float32) std np.array([0.229, 0.224, 0.225], dtypenp.float32) image (image - mean) / std # NHWC - NCHW image torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0) return image后处理就是取 logits 的 argmax然后映射到预测颜色。这个映射关系在训练时最好固定下来方便对比可视化结果def postprocess(logits: torch.Tensor, color_map: np.ndarray) - np.ndarray: pred torch.argmax(logits, dim1).squeeze(0).cpu().numpy() # color_map 是一个 NUM_CLASSES x 3 的 RGB 数组 vis color_map[pred] return vis预处理和后处理是整个推理链路里最容易出错也最不容易被发现的部分建议你在跑任何模型之前先用单张图片验证这两段代码的输入输出形状。5. 训练阶段的实战配置与踩坑记录5.1 损失函数和类别权重语义分割最常用的损失还是交叉熵。但直接平铺交叉熵在一个类别不平衡的数据集里非常吃亏。比如遥感地块分割背景像素占了 80%建筑物只占 5%网络学到全预测背景也能有 80% 的准确率但 MIOU 惨不忍睹。我的方案是两类处理一是 use weighted cross entropy根据每个类别在训练集中出现的像素频率计算权重。计算方式是weight_i median_freq / freq_ifreq 越小的类权重越大。二是叠加一个 Dice Loss直接优化分割区域和真实区域的像素交集import torch.nn as nn import torch.nn.functional as F class DiceLoss(nn.Module): def __init__(self, smooth1.0): super().__init__() self.smooth smooth def forward(self, logits, target): probs F.softmax(logits, dim1) target_onehot F.one_hot(target, num_classeslogits.size(1)).permute(0, 3, 1, 2).float() intersection (probs * target_onehot).sum(dim(2, 3)) union probs.sum(dim(2, 3)) target_onehot.sum(dim(2, 3)) dice (2.0 * intersection self.smooth) / (union self.smooth) return 1.0 - dice.mean()总损失我设置为weighted_ce 0.3 * dice_loss。这个 0.3 的系数是我经过几组实验选出来的经验值太小了 Dice Loss 不起作用太大了会干扰交叉熵的收敛方向。5.2 优化器、学习率与训练轮数优化器我推荐直接用 AdamWPyTorch 内置的torch.optim.AdamW就行。weight decay 设成 1e-4正值能抑制轻量网络的过拟合。学习率策略用的是 warmup cosine decay。前 5 个 epoch 从 3e-4 线性升到 3e-3之后按余弦曲线衰减到 3e-5。这个组合在小模型上非常好用前期稳定性明显提高。训练轮数方面我的经验是 30-50 epoch 基本就能看出模型的真实水平。如果 30 epoch 结束 MIOU 还在持续上升就继续加 20 轮如果 loss 已经平稳就没必要浪费算力。Batch size 看显存。11GB 显存的单卡512x512 输入batch size 能开到 8。如果你只有 6GB 显存batch size 降到 4 也不会对结果产生毁灭性影响但要记得把编码器前几个 stage 的 BatchNorm 冻结住否则 BN 统计量在小 batch 下会非常不稳定。5.3 数据增强策略分割任务的数据增强不能太激进因为每个像素的标签必须跟着图像做同步变换。我做的是下面四件套随机水平翻转概率 0.5随机缩放 0.5 到 1.5 倍然后随机裁剪到 512x512颜色抖动亮度、对比度、饱和度抖动随机旋转 90 度的倍数需要特别注意的是随机旋转 90 度这种增强如果数据集里的物体有明显方向性比如车道线或者建筑物立面旋转增强可能引入不合理的样本。我一般只在室内场景数据集上开启这个增强遥感地物数据集就关掉。5.4 我踩过的三个坑第一个坑是 BatchNorm 统计量在冻结阶段被意外更新。有些同学把 encoder 整体requires_gradFalse之后以为 BN 就不会变了。但 PyTorch 里 BN 的 running_mean 和 running_var 更新不依赖 requires_grad只要网络还在训练模式BN 就会继续用当前 batch 更新统计量。我当时的处理是单独把前两个 stage 的 BN 设成 eval 模式def freeze_bn(m): if isinstance(m, nn.BatchNorm2d): m.eval() self.encoder.stage1.apply(freeze_bn) self.encoder.stage2.apply(freeze_bn)第二个坑是混合精度训练导致的分割边界抖动。AMP 在图像分类上几乎无感但分割任务对逐像素的细节敏感。我在某些 batch 上发现预测图出现了一小片一小片的噪点排查下来是 AMP 的 grad scaler 在 loss 较小时精度损失过大。解决方式是只对 decoder 部分用 AMPencoder 保持 FP32或者干脆全程 FP32反正 MobileNetV1 的计算量不大。第三个坑是数据加载瓶颈。MobileNetV1 前向推理非常快如果数据加载不做多进程和预取GPU 利用率经常掉到 60% 以下。我用了torch.utils.data.DataLoader的num_workers8和prefetch_factor4线上训练速度直接提升将近一倍。6. 结果评估与项目打包如何判断分割效果是否真的达标6.1 评估指标和可视化语义分割最常用的指标是 MIOU这个指标衡量预测结果和真实标签在每个类别上的交并比平均值。我先说清楚计算方法逐类计算(预测区域 ∩ 标签区域) / (预测区域 ∪ 标签区域)然后对所有类别取平均。MIOU 比全局准确率要严格得多它不会让小类别被大类别完全掩盖。我评估时还看两个辅助指标平均像素准确率 PA 和频率加权交并比 FWIoU。PA 会偏向大类别FWIoU 会给高频率类别更高权重这两个指标和 MIOU 一起看能更全面判断模型在哪些类别上偏弱。举一个实测例子。我在一个四类室内场景数据集上训练了这个 MobileNetV1-U-Net输入 512x51230 个 epoch 后 MIOU 大约 67%。其中地面类别 MIOU 最高接近 85%墙面次之而小物体类别只有 42%。这说明模型对大面积的背景类学习得比较好对小目标的细节还需要更多数据或更强的解码器。除了数值指标可视化也很重要。我会保存三种图原始图像、标注 mask、预测 mask并排成一张三联图。看可视化时重点观察两点一是物体边界是否整齐二是小目标是否完整。这两个问题在数值上未必能直观反映但人眼一看就很明显。6.2 TorchScript 导出和项目打包训练完成的模型不能只在 Python 脚本里跑整套项目要交付出去就必须考虑模型导出。PyTorch 最常见的导出格式是 TorchScript。导出时需要注意模型里的nn.Upsample在 TorchScript 下是支持的但如果你想用torch.jit.trace而不是script就必须把scale_factor参数固定。用size参数指定目标尺寸会比scale_factor更稳定import torch model MobileNetV1UNet(num_classes4) model.load_state_dict(torch.load(best_mobilenet_unet.pth)) model.eval() dummy_input torch.randn(1, 3, 512, 512) scripted torch.jit.trace(model, dummy_input) scripted.save(mobilenet_unet_traced.pt)导出前记得把模型切到 eval 模式否则 BN 层的 running_mean 还在更新导出的模型行为会和训练时不一样。打包整套项目的时候我的习惯是必须具备以下几样东西README写明环境安装步骤、数据格式、训练和推理命令requirements.txt固定 PyTorch、torchvision 等核心依赖版本预训练权重文件一份结构完整的 config.py一个单图预测的演示脚本和测试图片这套项目的全套意义就在这里别人拿到手不需要跟你反复确认就能从头跑起来。6.3 后续改进方向MobileNetV1-U-Net 只是轻量化分割的起点。如果你后续想把精度再往上提可以尝试三个方向第一把编码器换成 MobileNetV2 或 MobileNetV3。它们在相同算力下精度更高尤其是 V3 的 SE 模块对通道注意力建模有明显帮助但部署时要确认推理引擎对 h-swish 的支持。第二在编码器和解码器之间加一个轻量 ASPP 模块。这个模块用不同膨胀率的空洞卷积并行捕捉多尺度上下文能有效提升不同大小目标的分割准确率增加的参数量也不多。第三把损失函数换成 Lovasz-Softmax 或 Focal Loss在小目标和类别不平衡场景下可能带来更可观的提升。我自己在实际训练中的体会是轻量化模型不是不能出精度只是要把每一处细节都扣到位。编码器选择、解码器通道数、损失函数加权、学习率安排任何一个环节偷懒最终 MIOU 都会给你颜色看。这套 MobileNetV1-U-Net 项目跑通之后我在自己的数据集上反复迭代了好几轮最终固定下来的这套配置已经可以作为一个稳定的基线版本后续不管换骨干还是加模块都能以它为参照来对比收益。