做语义分割项目,很多人第一反应就是上DeepLabV3+或者SegFormer这类大模型,其实在工程落地和学术baseline场景下,MobileNetV1-Unet的组合远比想象中实用。最近我把整套基于PyTorch的MobileNetV1-Unet图像分割项目完整梳理了一遍,从网络搭建、数据准备、训练调参到评估推理全部跑通,踩了不少坑也沉淀了不少可复用的经验。这篇就把整个项目的技术方案和实操细节完整拆开讲清楚,适合刚入门语义分割、想快速搭建一个完整baseline的同学,也适合需要在移动端或者在有限算力条件下跑分割任务的开发者参考。
1. 整体方案设计与网络选型思路
1.1 为什么选MobileNetV1作为编码器
语义分割本质上是对图像进行逐像素的密集分类,编码器负责提取特征,解码器负责恢复空间分辨率。编码器的选择直接决定了整个网络的运算量和分割精度。我在这个项目里选MobileNetV1做主干的理由很直接:它是深度可分离卷积的代表作,把标准3x3卷积拆成depthwise卷积加pointwise卷积,计算量降到标准卷积的九分之一左右,同时参数量也显著减少。
用一句生活化的话来理解:标准卷积就像让十几个工人同时干同一块地的活,大家都要弯下腰去搬砖,每个动作都重复做;而深度可分离卷积则是先让每个工人只处理自己面前的一个通道,干完后再由另一个小组把各通道的结果混合汇总。干同样多的活,前者每个环节都要动用全部人手,后者每个环节只动用一部分人手,整体开销自然低很多。
在ImageNet上预训练过的MobileNetV1权重重用价值很高,作为encoder做迁移学习效果很稳定。相比ResNet-34、ResNet-50这类编码器,MobileNetV1在GPU显存占用和推理速度上的优势非常明显。我用同一批数据分别测过ResNet34-Unet和MobileNetV1-Unet,在batch size为8、720x720输入分辨率下,ResNet34跑一个epoch耗时约为MobileNetV1的1.67倍,而mIoU差距基本能控制在1.5个百分点以内。对于绝大多数需要快速迭代的训练场景,MobileNetV1已经是效率和精度平衡得很好的选择。
1.2 解码器与跳跃连接的设计考量
UNet的核心卖点就是跳跃连接。编码器在下采样过程中逐层丢失空间细节,解码器通过skip connection把对应层的特征图拼回来,让网络在恢复分辨率的同时也能利用编码器中间层保留的细节信息。我采用的是UNet标准设计:四个stage的特征图分别经过1x1卷积调整通道数,然后与解码器的上采样结果在通道维度上拼接,再经过两个3x3卷积块融合。
这里有几个细节值得注意。第一,MobileNetV1的每个stage输出的channel数是固定的,与UNet解码器期望的channel不完全匹配,所以跳跃连接前必须加一个1x1卷积做通道对齐。第二,拼接之后接的双卷积块我建议还是用BatchNorm加ReLU,而且BatchNorm在推理阶段必须走整体统计量,否则预测时会出现明显的块状噪声。第三,MobileNetV1最后一个特征图分辨率是输入的1/32,这个低层特征语义信息最丰富,解码器第一层直接对它上采样,同时在深层编码器特征与解码器特征拼接时,我会额外加一个空间注意力门控,在不引入太多参数的前提下进一步抑制背景响应。
我实测过有没有跳跃连接的区别:不加skip connection的版本在边缘区域的分割质量明显下降,尤其是一些细小物体孔洞区域,mIoU掉了大概9个百分点。所以UNet架构里的跳跃连接不只是一个特色,而是切切实实的精度保障器。
1.3 为什么用PyTorch而不是其他框架
PyTorch在这种项目中的优势是动态计算图和灵活的forward控制流。语义分割的解码过程每个stage结构相似但又需要显式的上采样操作,PyTorch的nn.Upsample、F.interpolate或者nn.ConvTranspose2d都很好用,而且改起来直观。加上torchvision里直接集成了MobileNetV1的预训练权重,torch.load就可以直接加载。
另一个让我坚持PyTorch的原因是其生态对图像分割极其友好。albumentations库和PyTorch Dataset的配合几乎是一体化的,多进程数据加载配合pin_memory能大幅减少GPU等待时间。调试阶段直接用torchviz或者TensorBoard的add_graph可以可视化网络结构,排查维度问题很方便。分布式训练方面,torchrun和DDP也足够应付多卡场景。
2. 网络结构解析与核心代码实现
2.1 MobileNetV1编码器各层输出尺寸
MobileNetV1编码器以标准3x3卷积加BatchNorm加ReLU开始,然后串联深度可分离卷积下采样块。输入图片如果是3x720x720,各stage输出的特征图尺寸如下表所示:
| 层名 | 操作方式 | 输出尺寸 | 通道数 |
|---|---|---|---|
| Stage0 | 标准卷积stride=2 | 360x360 | 32 |
| Stage1 | 深度可分离卷积stride=2 | 180x180 | 64 |
| Stage2 | 深度可分离卷积stride=2 | 90x90 | 128 |
| Stage3 | 深度可分离卷积stride=2 | 45x45 | 256 |
| Stage4 | 深度可分离卷积stride=2 | 23x23 | 512 |
| Stage5 | 深度可分离卷积stride=2 | 12x12 | 1024 |
其中Stage1到Stage4之间还有若干个不改变分辨率的深度可分离卷积层,它们在UNet里负责提供语义信息更丰富的中间特征。我通常在Stage2到Stage5这四个位置取特征图做跳跃连接,分别对应解码器四个阶段的输入。需要注意的是,原始MobileNetV1分类模型末尾还有全局平均池化和全连接层,在用于分割编码器时要全部去掉,只需要保留到最后一个卷积特征输出。
2.2 UNet解码器结构与拼接流程
解码器整体是一个四级金字塔结构,从Stage5输出开始逐级恢复分辨率。每一级解码块的组成是:上采样将特征图尺寸扩大一倍(我选用双线性插值配合3x3卷积,而不是转置卷积,因为前者不容易产生棋盘格伪影),然后与对应编码器特征图做通道维度拼接,再通过连续两个3x3卷积块融合特征。最终输出的分割头通过1x1卷积映射到类别数。
拼接流程用伪代码表示:
# 假设encoder输出为c5, c4, c3, c2,依次对应Stage5到Stage2 d4 = decoder_stage(c5) # 12x12 -> 23x23 d4 = upsample(d4) # 23x23 -> 45x45 d4 = conv_block(cat([d4, c4])) # 拼接Stage3特征(45x45) d3 = upsample(d4) # 45x45 -> 90x90 d3 = conv_block(cat([d3, c3])) # 拼接Stage2特征(90x90) d2 = upsample(d3) # 90x90 -> 180x180 d2 = conv_block(cat([d2, c2])) # 拼接Stage1特征(180x180) d1 = upsample(d2) # 180x180 -> 360x360 d1 = conv_block(cat([d1, c1])) # 拼接Stage0特征(360x360) out = seg_head(d1) # 1x1卷积输出N类logits整个解码器结构比较规整,关键是每一级的输入输出通道数要仔细核对。我在最初搭结构时曾把Stage4和Stage5的顺序搞反了,导致拼接时特征图尺寸不一致直接报错,后来在代码里加了assert断言:上一级上采样后shape必须与跳跃连接特征完全一致。
2.3 深度可分离卷积的PyTorch实现要点
深度可分离卷积的PyTorch实现非常简单,核心就是nn.Conv2d的两个参数groups设置。depthwise卷积设置groups=in_channels,每个输入通道分别做卷积,pointwise卷积使用kernel_size=1将通道混合。下面是标准的MobileNet卷积块定义:
class DepthwiseSeparableConv(nn.Module): def __init__(self, in_channels, out_channels, stride=1): super().__init__() self.depthwise = nn.Conv2d( in_channels, in_channels, kernel_size=3, stride=stride, padding=1, groups=in_channels, bias=False ) self.pointwise = nn.Conv2d( in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=False ) self.bn1 = nn.BatchNorm2d(in_channels) self.bn2 = nn.BatchNorm2d(out_channels) self.relu = nn.ReLU(inplace=True) def forward(self, x): x = self.depthwise(x) x = self.bn1(x) x = self.relu(x) x = self.pointwise(x) x = self.bn2(x) x = self.relu(x) return x这里有个关键点:depthwise卷积的计算量虽然少,但它逐通道处理会丢失通道间的信息交互,pointwise卷积就是负责各通道线性组合的。两个操作中间必须要有BatchNorm和ReLU,直接把两个卷积串在一起会导致训练时梯度传播不稳定。另外我建议depthwise卷积不使用bias,因为后面接BatchNorm会让bias冗余,加了反而增加过拟合风险。
2.4 预训练权重加载与模型微调
MobileNetV1的预训练权重在torchvision中可以直接获取。加载时需要注意权重key与自定义模型的一致性。我为了方便将encoder包在一个名为MobileNetEncoder的Module里,加载时使用带encoder.前缀匹配的方式:
import torchvision.models as models def build_model(num_classes, pretrained=True): encoder = MobileNetEncoder(pretrained=pretrained) decoder = UNetDecoder(num_classes) model = MobileNetUNet(encoder, decoder) if pretrained: state_dict = torchvision.models.mobilenet_v1(weights=models.MobileNet_V1_Weights.IMAGENET1K_V1).state_dict() # 过滤掉分类头权重,只保留卷积层 filtered = {k: v for k, v in state_dict.items() if 'classifier' not in k} model.encoder.load_state_dict(filtered, strict=False) return model严格模式建议设成strict=False,因为MobileNetV1的预训练模型里有些层的命名和encoder子模块不完全一致,完全严格加载反而容易报错。加载完预训练权重后,我一般会额外打印一句提示:是否所有encoder层权重都被正确初始化。如果有些层是随机初始化,训练时这些层学习速度会明显更慢,影响整体收敛。
3. 数据准备与预处理全流程
3.1 数据标注格式与目录结构
不管做哪类分割任务,数据准备决定了模型性能的上限。语义分割需要的标注一般是和输入图像同尺寸的单通道PNG,像素值就是类别编号。比如背景为0、道路为1、车辆为2。一定要避免把类别从1开始编号,因为很多损失函数和评估指标默认0是背景,混乱编号会导致mIoU统计出错。
我推荐的目录结构如下:
data/ ├── images/ │ ├── train/ │ │ ├── 00001.jpg │ │ └── ... │ ├── val/ │ └── test/ ├── masks/ │ ├── train/ │ │ ├── 00001.png │ │ └── ... │ ├── val/ │ └── test/ └── class_dict.jsonclass_dict.json文件维护类别和颜色的映射关系,方便可视化。比如{"background": 0, "road": 1, "building": 2}。可视化时把mask像素值映射到RGB颜色,再与原图叠加,能直观判断标注质量。
3.2 数据增强策略与实现
数据增强在语义分割里一定要谨慎处理。图像和mask必须用完全相同的几何变换,否则模型会学到错误的空间对应关系。我使用albumentations库定义增强pipeline,它支持以字典方式同时处理image和mask:
import albumentations as A from albumentations.pytorch import ToTensorV2 def get_train_transforms(): return A.Compose([ A.RandomResizedCrop(height=360, width=360, scale=(0.8, 1.0), p=0.5), A.HorizontalFlip(p=0.5), A.RandomRotate90(p=0.5), A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.1, rotate_limit=10, border_mode=0, value=0, mask_value=0, p=0.5), A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.3), A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ToTensorV2(), ], additional_targets={'mask': 'image'}) def get_val_transforms(): return A.Compose([ A.Resize(360, 360), A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ToTensorV2(), ], additional_targets={'mask': 'image'})这里注意几个细节:ShiftScaleRotate里border_mode=0表示填充的边界像素值为0,value=0对应图像边界填充值,mask_value=0是mask边界填充值。因为mask里0通常是背景,所以填充0不会引入新类别。RandomResizedCrop的scale参数如果太低会让物体被过度裁剪,导致标签碎片化,这个参数我建议保持在0.8以上。Normalize的均值和方差必须和ImageNet保持一致,因为预训练权重的特征分布是在这个归一化条件下学出来的。
3.3 Dataset类与数据加载加速
自定义Dataset类要干三件事:读图像、读对应mask、应用增强。PyTorch的标准Dataset接口就是__len__和__getitem__,写成下面的形式比较方便:
class SegmentationDataset(Dataset): def __init__(self, img_dir, mask_dir, transforms=None): self.img_paths = sorted(glob.glob(os.path.join(img_dir, '*.jpg'))) self.mask_paths = sorted(glob.glob(os.path.join(mask_dir, '*.png'))) self.transforms = transforms assert len(self.img_paths) == len(self.mask_paths) def __len__(self): return len(self.img_paths) def __getitem__(self, idx): image = cv2.imread(self.img_paths[idx]) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) mask = cv2.imread(self.mask_paths[idx], cv2.IMREAD_GRAYSCALE) if self.transforms: transformed = self.transforms(image=image, mask=mask) image = transformed['image'] mask = transformed['mask'] mask = mask.long() return image, mask读图用cv2而不是PIL,因为cv2的imread速度比PIL快很多,多进程数据加载时能减少CPU瓶颈。mask要转成长整型,因为CrossEntropyLoss要求target是整型索引而不是one-hot。增加num_workers数量能显著加快数据供给,但也要根据CPU核数设置,一般设置为CPU核数的一半到全部之间。另外把DataLoader的pin_memory=True开启,GPU训练时数据从固定内存拷贝到显存的速度会更快。
4. 训练配置与调参经验
4.1 损失函数选择与实现
语义分割常见损失函数有CrossEntropyLoss、Dice Loss和Focal Loss。类别分布不均衡时,直接使用CrossEntropyLoss容易让模型偏向多数类。我这个项目里类别分布严重不均衡,有些类别只占整体像素的1%左右,所以使用了CrossEntropyLoss和Dice Loss的加权组合。
其中CrossEntropyLoss负责逐像素的精细分类,Dice Loss负责整体的区域重叠度优化,两者的加权比例我常用0.6比0.4。Dice Loss的PyTorch实现不算复杂,这里给出一个稳定的版本:
def diceloss(pred, target, smooth=1.0): num_classes = pred.size(1) pred_softmax = F.softmax(pred, dim=1) target_onehot = F.one_hot(target, num_classes).permute(0, 3, 1, 2).float() intersection = (pred_softmax * target_onehot).sum(dim=(2, 3)) cardinality = pred_softmax.sum(dim=(2, 3)) + target_onehot.sum(dim=(2, 3)) dice = (2.0 * intersection + smooth) / (cardinality + smooth) return 1.0 - dice.mean()注意类别many类别的Dice分数会对整体产生比较大的拉动,所以在损失里我做了一个权重衰减:对每个类别计算Dice后乘以类别频率的反比,再取均值。这个方法比简单取平均更稳,尤其是在目标类别非常小的情况下。另外,smooth参数(平滑系数)很重要,把它设成1.0而不是极小的值,可以避免在训练初期某些类别的GT和预测完全没有交集时导致梯度爆炸。
如果场景是单类分割,也就是只有前景和背景,可以把上面的实现进一步简化,直接用sigmoid加BCE实现类似的效果。多类别场景还是用上面的softmax版本更合适。
4.2 优化器、学习率与训练轮次配置
我使用AdamW作为优化器,初始学习率设置为3e-4。相比SGD,AdamW在分割任务里收得更快、更稳,weight decay我设为1e-4,防止深层特征过拟合。配合CosineAnnealingLR学习率调度器,第一个epoch经历一个小的warmup,从1e-5逐渐升到3e-4,之后余弦衰减到训练结束。总训练轮次我设为120个epoch,batch size视显存大小而定,720x720输入下单卡12GB大概能跑batch size为8。
其实作为第一次跑这个项目的建议:先用一个很小的子集跑3到5个epoch验证损失能否下降、显存是否够用,再上全量数据。我最初直接全量训练,结果到第10个epoch发现loss不降,排查了一圈发现是DataLoader的预处理在mask上多做了一次归一化,导致target被转成了float,CrossEntropyLoss直接报错,但报错信息不够直观,浪费了大半天。
训练循环的核心代码可以这样组织:
for epoch in range(start_epoch, num_epochs): model.train() train_loss = 0.0 for images, masks in train_loader: images, masks = images.cuda(), masks.cuda() outputs = model(images) loss = criterion(outputs, masks) optimizer.zero_grad() loss.backward() optimizer.step() train_loss += loss.item() scheduler.step() val_loss, val_miou = validate(model, val_loader) save_checkpoint(model, optimizer, epoch, val_miou)这里有一个很容易踩的坑:model.train()和model.eval()切换时,BatchNorm的running_mean和running_var更新行为不同。如果验证阶段忘记切换,BatchNorm一直用mini-batch统计量,预测结果会非常不稳定,分割图会出现一块一块的伪影。这个问题我以前在实时推理时亲身体验过,最后查了很久才发现是BN的坑。
4.3 编码器冻结与解冻策略
迁移学习的一个常见做法是训练初期冻结编码器,只训练解码器。这能加快收敛,也会减少显存占用,因为冻结层不需要计算梯度。我建议先用5个epoch冻结编码器训练解码器,让解码器先学会利用MobileNetV1提取的特征;之后解冻编码器,用较低的学习率(比如编码器学习率为解码器的十分之一)继续联合训练,这样对预训练特征的破坏最小。
实现时很简单,给不同参数组设置不同的学习率:
encoder_params = [] decoder_params = [] for name, param in model.named_parameters(): if 'encoder' in name: encoder_params.append(param) else: decoder_params.append(param) optimizer = torch.optim.AdamW([ {'params': decoder_params, 'lr': 3e-4}, {'params': encoder_params, 'lr': 3e-5}, ], weight_decay=1e-4)4.4 训练过程的监控与日志记录
训练过程最忌讳盲训练。我通常在每个epoch结束后在验证集上计算mIoU、Dice、F1和像素准确率,同时保存最佳模型。历史最佳mIoU的checkpoint单独存一份,当前epoch的checkpoint存成另一份,避免断电或中断导致训练白跑。记录指标时用TensorBoard或者wandb都行,主要是看loss曲线和验证集指标曲线。我发现一个规律:验证集的mIoU曲线如果在前20个epoch内涨幅很小,大概率是学习率初始值太大,导致loss在震荡;如果loss在初期下降得很快但mIoU停滞,可能是数据增强太强,或者验证集和训练集的分布差异过大。
5. 评估指标与推理部署细节
5.1 mIoU计算逻辑与实现
语义分割的评估以mIoU为核心指标。mIoU的计算方式是对每个类别分别计算IoU,然后取平均。单类IoU的计算是交集除以并集:预测和GT共同的像素数除以预测和GT合集的像素数。实现时可以先构造混淆矩阵,再逐类计算:
def compute_miou(pred, target, num_classes): ious = [] pred = pred.view(-1) target = target.view(-1) for cls in range(num_classes): pred_inds = (pred == cls) target_inds = (target == cls) intersection = (pred_inds & target_inds).sum().float() union = (pred_inds | target_inds).sum().float() if union == 0: ious.append(float('nan')) else: ious.append((intersection / union).item()) return np.nanmean(ious)注意对GT和预测中都不出现的类别,union为0,这个类别不算在平均值里,也就是跳过而不是当作0。如果当作0处理,当类别数量比较多时,mIoU会被严重拉低,这会对模型好坏产生错误判断。
5.2 推理流程与后处理技巧
训练完成后,推理阶段通常要处理任意尺寸的输入图。如果图片太大,直接放进去会显存溢出,我的做法是用滑窗切块推理,然后把结果拼接回原始尺寸。滑动窗口重叠区域如何处理是个细节问题,我用的是重叠部分取均值或者取最大值,经验是取均值更平滑,取最大值对边界更锐利,但容易在重叠边界处出现轻微的不连续感。实际项目里重叠区域设25%,取平均即可。
预测结果的可视化也很重要,直接把argmax结果的索引图存成一个灰蒙蒙的图看不出效果。需要预先准备每种类别对应的颜色映射表,生成一张RGB结果图,再与原图按一定透明度叠加。比如对于道路类别标成红色、车辆标成蓝色、建筑标成绿色,叠加后的效果图更适合直接汇报查看。
5.3 模型导出与加速方案
如果项目需要实际部署,可以把PyTorch模型导出成TorchScript或ONNX格式。MobileNetV1-Unet这种结构导出ONNX很顺畅,只要注意在forward里不要用Python的原生if分支和循环,ONNX导出时这些动态控制流会麻烦一些。我的forward流程全是卷积和拼接操作,所以导出非常顺利。导出ONNX后可以进一步用TensorRT优化推理速度,在NVIDIA GPU上实测可以做到3倍左右的提速。
如果部署目标是CPU或者移动端,可以考虑量化。PyTorch官方的量化工具支持对称量化和非对称量化。MobileNetV1本来就是为了轻量化设计的,量化后模型大小只有原来的四分之一左右,推理速度也能提升。不过量化对分割模型精度有一定影响,通常mIoU会掉几个点,需要评估量化前后指标来决定是否可用。
6. 项目代码组织与复用建议
6.1 项目目录结构与模块划分
一套完整的分割项目,代码结构一般分成数据、模型、工具、配置、训练入口和推理入口几个部分。下面是我习惯的目录组织方式:
project/ ├── config.py # 全局配置参数 ├── models/ │ ├── mobilenet_encoder.py │ ├── unet_decoder.py │ └── mobilenet_unet.py ├── data/ │ ├── dataset.py # Dataset类与数据增强 │ └── __init__.py ├── utils/ │ ├── metrics.py # mIoU/Dice等评估函数 │ ├── loss.py # 损失函数 │ ├── visualize.py # 可视化工具 │ └── checkpoints.py # 模型保存与加载 ├── train.py # 训练脚本 ├── validate.py # 验证脚本 ├── inference.py # 推理脚本 └── export_onnx.py # 模型导出脚本6.2 复用到其他分割场景的改造方法
这个网络不只能做某一种固定分割任务。如果换成其他类别的分割,只需要修改两个地方:一是数据集中mask的类别数和Class dict;二是模型最终的输出通道数(num_classes)。比如医学图像分割、遥感图像地物分类、工业质检表面缺陷检测,结构上完全可以复用。数据增强部分根据场景适当调整,像医学图像可以做弹性形变增强,遥感图像可以加随机旋转和翻转,工业检测一般保持原尺寸不变、增加光照扰动。
6.3 依赖环境与版本兼容性
PyTorch版本建议使用不低于1.10的版本,2.x版本也可以正常使用,但是有些旧代码里用到的torchvision.models.mobilenet_v1权重参数在不同版本下有差异。我建议固定使用某个已经验证过的组合,比如Python 3.9、PyTorch 1.13.1加torchvision 0.14.1、CUDA 11.7,这个组合在绝大多数环境下都能稳定运行。albumentations库版本不要低于1.3,因为低版本在字典数据的处理上有些API不兼容。
7. 常见问题与排查技巧实录
7.1 显存溢出问题
显存溢出是最常见的问题。如果训练在某个batch就OOM了,优先做这几件事:调小batch size、降低输入分辨率、检查是否解冻了编码器。其中解冻编码器会让反向传播需要保存更多中间激活值,显存占用会显著上升。如果单卡显存不够又不打算换卡,可以考虑用梯度累积模拟更大的batch size,但要注意BatchNorm的行为,因为BN统计的是当前mini-batch的均值和方差,梯度累积时BN仍然按小batch更新,可能带来精度损失。
我实测在12GB显存的GPU上,720x720输入、batch size 8是MobileNetV1-Unet的稳定工作点。如果显存只有8GB,可以把输入降到576x576或者batch size降到4。不要把输入直接降到256x256,那样分割结果的边界会很毛糙,因为上采样恢复的空间细节不够。
7.2 Loss长时间不下降的排查清单
Loss不下降先别急着调模型,按下面的清单排查:
- 检查数据加载:确认图像和mask是否一一对应,多图少图都会导致错位。我遇到过mask文件多了几个,且文件名对不上,导致相近样本错位训练的情况。
- 检查归一化:确认输入图像做了ImageNet归一化,mask没有做归一化,仍然保持整数索引。
- 检查标签范围:确认mask中最大像素值不超过num_classes-1,如果有超出,CrossEntropyLoss会直接报错或者数值异常。
- 检查学习率:初始学习率太大会使loss震荡或者弥散,太小的又会让loss下降极慢,建议先跑小数据测试收敛速度。
- 检查类别平衡:如果某些类别只有极少像素,模型会把该类预测成背景。解决方法是加权重或者在损失函数里对稀有类提高权重。
7.3 分割结果出现块状伪影
分割图出现块状伪影的常见原因有三个。第一个是BatchNorm在推理阶段没有切换到eval模式,或者使用的是InstanceNorm,这是因为在train/eval切换时忘记设置model.eval()。第二个是数据增强时对图像和mask用了不一样的变换,导致局部区域出现空间错位。第三个是上采样用的转置卷积kernel_size和stride设置不匹配,产生棋盘格效应,业内通常叫checkerboard artifact。解决方式是改用双线性插值上采样,或者对kernel_size取整倍关系并配合weight initialization。
7.4 类别不均衡的实战处理
语义分割里类别不均衡是家常便饭,有时目标类别只占整体像素的千分之一。除了loss里加权,我还有一个比较实用的技巧:对图像做失焦裁剪。比如训练时以一定的概率只裁剪原图中目标类别占比高的区域,这样相当于隐式地把稀有类别的采样比例提高了。这个方法处理小目标检测非常有效,而且不用改任何模型结构。还有一个传统做法是当模型训练到后期时,对分类错误的像素做hard example mining,加重这些区域在loss中的比重,往往能再提升1到2个点mIoU。
8. 最终使用体会
做这个项目整体的体会是:MobileNetV1-Unet架构把效率和精度的平衡做得非常出色,在标注数据量不够大、GPU不算多的情况下,它几乎是最合适的起步方案之一。我最初用ResNet34编码器跑同样的数据,精度高一点点,但训练时间几乎翻倍,显存占用也高不少,对于快速验证想法非常不友好。后来切到MobileNetV1,一版训练从原本的四天缩到两天半,迭代效率高了很多。
我个人实际跑完这个项目后留下一个习惯:所有语义分割模型的第一个版本一律先按这个结构搭,因为数据的坑远远比模型的坑多,先用一个稳定的baseline把数据、增强、评估链路验证好,再考虑换更强的主干或者引入注意力机制。很多同学一上来就照搬大模型,结果被显存、训练时间、调参复杂度拖得寸步难行,反而拿不出一个可靠的结果。
最后再分享一个小技巧:MobileNetV1-Unet训练早期可以把输入分辨率设低一些跑通链路,等确认一切正常再调高分辨率正式训练,这样可以省下不少试错时间。模型本身很轻,加上预训练权重,几乎是从零搭建分割baseline成本最低的一条路。希望这篇完整细致的拆解能帮你少走弯路,也欢迎在实际复现中和我交流各自的调参心得。