ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

BiSeNet人脸解析19类分割:从PyTorch训练到端侧部署全流程实战

BiSeNet人脸解析19类分割:从PyTorch训练到端侧部署全流程实战 1. 人脸解析到底在做什么从BiSeNet的19类分割说起人脸解析Face Parsing这个词听起来挺学术但说白了就是给一张人脸照片里的每个像素贴标签——这块是左眉毛那块是右眼珠嘴唇归嘴唇头发归头发。它和普通人脸检测最大的区别在于检测只告诉你“脸在哪”解析要告诉你“脸的每个部分分别在哪”。我第一次接触这个任务的时候以为跟语义分割差不多后来发现细节上的坑完全不是一回事。BiSeNetBilateral Segmentation Network是这个人脸解析任务里最常被拿来做落地的骨干网络之一。它的核心思路是双路结构一条Spatial Path保留高分辨率特征专门管边缘和细节另一条Context Path用轻量级骨干快速下采样负责抓全局语义。最后用一个Feature Fusion Module把两路拼起来。这个设计的好处是推理速度快同时在边缘区域不会糊成一团。对于人脸解析这种对五官边界极其敏感的任务BiSeNet比那些单纯堆深度的网络要实用得多。19类语义分割是这套方案的标准输出配置。具体类别通常包括背景、脸部皮肤、左眉、右眉、左眼、右眼、左耳、右耳、鼻子、上唇、下唇、嘴部内部、上牙、下牙、舌头、脖子、头发、眼镜、以及配饰等。不同数据集和预训练权重会有细微差异但大体上就是把人脸拆成这些语义区域。这个粒度足够支撑绝大多数应用虚拟试妆需要精确的唇部和眼部区域人脸编辑需要头发和皮肤的分离表情迁移需要五官的独立控制。我之所以选这个题目来写是因为过去半年里我在三个不同项目里都用了BiSeNet做人脸解析从PyTorch训练到ONNX导出再到TFLite和NCNN的端侧部署几乎把整条链路踩了一遍。这篇文章会把整个流程拆开讲清楚包括模型选型、数据准备、训练调参、格式转换、量化压缩、端侧推理以及那些文档里不会写的坑。不管你是刚接触语义分割的新手还是已经在做端侧部署的老手应该都能从里面找到能直接用的东西。2. 为什么选BiSeNet而不是其他分割网络2.1 人脸解析对网络结构的特殊要求人脸解析和通用语义分割有一个本质区别它的目标区域非常精细而且区域之间的边界往往只差几个像素。比如上唇和下唇之间的缝隙左眼和右眼之间的鼻梁区域这些地方的分类错误在视觉上会非常明显。通用的DeepLabV3或者PSPNet虽然mIoU指标好看但它们的输出步长通常是8或16意味着最终特征图的分辨率只有输入的1/8或1/16边缘细节丢失严重。BiSeNet的设计恰好针对这个问题。它的Spatial Path只有三层卷积输出步长是8但通过跳连和融合机制最终输出的特征图分辨率可以做到输入的1/4甚至1/2。这意味着每个像素的分类结果都能保留足够的空间信息。我在实际对比测试中发现同样训练到收敛BiSeNet在嘴唇和眉毛区域的IoU比DeepLabV3高出5到8个百分点而推理速度反而更快。另一个关键因素是参数量。人脸解析通常要部署在移动端或者边缘设备上模型不能太大。BiSeNet的完整版本参数量在49M左右如果换成MobileNetV3作为Context Path的骨干可以压到10M以内。这个量级对于ONNX和TFLite的端侧部署非常友好。2.2 BiSeNet与YOLO系列实例分割的区别最近YOLO26发布后很多人问实例分割和语义分割到底有什么区别能不能用YOLO直接做人脸解析。这里需要澄清一个概念实例分割输出的是每个独立目标的掩码比如“这是第一张脸”“这是第二张脸”语义分割输出的是每个像素的类别不区分个体。人脸解析属于语义分割因为一张脸就是一个整体不需要区分左右脸。YOLO系列的实例分割分支虽然也能输出掩码但它的掩码分辨率通常较低而且训练目标不同。用YOLO做人脸解析你会发现在五官边界上的精度远不如BiSeNet。我试过用YOLOv8-seg在同样的数据集上训练嘴唇区域的IoU只有BiSeNet的60%左右。所以如果你的任务是精细的人脸区域分割还是老老实实用BiSeNet这类专门为语义分割设计的网络。2.3 模型选型的实际考量在实际项目中我通常会准备两个版本的模型一个高精度的BiSeNet完整版用于服务端推理一个轻量化的BiSeNet-MobileNetV3版本用于端侧。服务端版本追求mIoU端侧版本追求速度和体积的平衡。这个策略的好处是你可以先用高精度模型验证业务逻辑等效果确认后再压缩部署。选型时还要考虑框架生态。BiSeNet在PyTorch上的实现最成熟预训练权重也最容易找到。如果你需要转ONNXPyTorch的导出工具链是最完善的。转TFLite的话需要先转ONNX再转TensorFlow中间会有一些算子兼容性问题后面会详细讲。3. 数据准备19类标注的坑与技巧3.1 数据集选择与标注规范公开的人脸解析数据集主要有CelebAMask-HQ、Helen、LFW-Parts等。CelebAMask-HQ是其中最常用的包含30000张高分辨率人脸图像标注了19个类别。但这个数据集也有问题部分图像的标注质量参差不齐尤其是头发和配饰区域有些标注明显是自动生成的边缘很粗糙。如果你要自己标注我建议用LabelMe或者CVAT。标注时要注意几个关键点第一类别定义要统一比如“嘴部内部”和“牙齿”的边界要明确第二边缘要尽量贴合实际轮廓不要为了省事画大块第三对于遮挡区域比如被头发遮住的耳朵要标注为不可见而不是强行猜测。我见过一个项目因为标注时把遮挡区域随便填了导致模型在推理时对遮挡区域产生严重误判。3.2 数据增强策略人脸解析的数据增强和通用分割不太一样。水平翻转是最有效的因为人脸基本对称翻转后类别也要对应交换左眼变右眼左眉变右眉。随机裁剪和缩放要小心因为人脸的关键区域比例不能变太多否则模型会学到错误的尺度关系。颜色抖动可以加但幅度不要太大。我试过用很强的颜色抖动结果模型对肤色变化过度敏感在深色皮肤上的表现明显下降。后来把亮度、对比度、饱和度的变化范围都控制在0.2以内效果就稳定了。还有一个容易被忽略的点训练时要保证每个batch里各个类别的样本均衡。人脸解析数据集中背景和皮肤通常占大部分像素而牙齿、舌头这些小类别样本很少。如果不做重采样或者加权损失模型会倾向于把所有像素都预测成背景和皮肤。我通常用类别加权的交叉熵损失权重根据每个类别的像素频率的倒数来设置。3.3 数据格式与预处理BiSeNet的输入通常是512x512的RGB图像标注是单通道的类别索引图。预处理时要注意归一化参数一般用ImageNet的均值和标准差。但人脸图像和ImageNet的分布有差异我实测下来用0.5的均值和0.5的标准差效果更好相当于把像素值归一化到-1到1之间。标注图在训练时要做one-hot编码或者直接用类别索引配合交叉熵损失。如果用Dice Loss或者Focal Loss需要先转成one-hot。我一般用交叉熵加Dice的混合损失交叉熵负责整体收敛Dice负责小类别的边界优化。4. PyTorch训练与调参实战4.1 训练环境搭建训练环境建议用PyTorch 1.12以上版本CUDA 11.6以上。BiSeNet的官方实现依赖比较多我建议直接用一个干净的conda环境然后手动安装依赖。主要需要的是torch、torchvision、opencv-python、numpy、tqdm、tensorboard。数据加载用DataLoadernum_workers设成4到8pin_memory打开。如果显存够大batch size可以设到16或32。我用RTX 3060 12G训练时batch size设16输入512x512显存占用大概8G左右。4.2 学习率与优化器选择BiSeNet训练我用的是SGD加动量初始学习率0.01动量0.9权重衰减1e-4。学习率调度用poly策略power设0.9。这个配置在CelebAMask-HQ上训练100个epoch左右能收敛到比较好的效果。也可以试试AdamW初始学习率3e-4权重衰减0.01。AdamW收敛更快但最终精度通常比SGD低一点点。如果训练时间紧用AdamW可以省不少时间。有一个细节BiSeNet的Context Path如果用了预训练的ResNet或MobileNet学习率要设小一点比如初始学习率的0.1倍。因为预训练权重已经很好太大的学习率会破坏原有的特征提取能力。4.3 损失函数设计前面提到我用的是交叉熵加Dice的混合损失。具体实现时交叉熵用nn.CrossEntropyLossweight参数传入类别权重。Dice Loss需要自己实现对每个类别计算Dice系数然后取平均。权重设置上背景和皮肤的权重设0.5其他类别设1.0到2.0。牙齿、舌头这些小类别可以设到3.0。但权重也不能太高否则模型会过度关注这些小类别导致整体mIoU下降。我一般会跑几组对比实验找到最优的权重配置。4.4 训练过程监控与调参训练时用TensorBoard监控loss和mIoU曲线。正常情况下loss在前10个epoch下降很快然后进入缓慢下降阶段。如果loss震荡严重可能是学习率太大或者batch size太小。如果loss下降很慢检查一下数据增强是不是太强了。验证集上的mIoU是最终指标。我一般每5个epoch验证一次保存mIoU最高的模型。如果验证集mIoU连续10个epoch不提升就降低学习率。如果降了两次还不提升就提前停止。还有一个经验训练后期可以关掉数据增强用原始数据微调几个epoch通常能提升0.5到1个百分点的mIoU。这个技巧在分割任务里很管用。5. 模型导出从PyTorch到ONNX的完整流程5.1 ONNX导出的关键参数PyTorch转ONNX用torch.onnx.export函数。关键参数有opset_version、input_names、output_names、dynamic_axes。opset_version建议用11或12兼容性最好。dynamic_axes可以把batch维度设成动态的方便后续部署时调整batch size。导出时要注意模型的输入输出格式。BiSeNet的输出通常是每个像素的类别概率形状是[batch, num_classes, height, width]。如果后续要做argmax得到类别索引可以在导出时加一个ArgMax节点也可以留给推理端处理。我一般留给推理端这样更灵活。导出后一定要用onnxruntime验证一下输出是否和PyTorch一致。我遇到过导出后输出形状对但数值不对的情况原因是某些算子在不同opset下的行为有差异。验证时用同一张输入图像分别跑PyTorch和ONNX比较输出的最大绝对误差应该小于1e-4。5.2 ONNX模型优化导出的ONNX模型可以用onnx-simplifier做简化去掉多余的算子合并常量节点。简化后的模型体积通常能减小10%到20%推理速度也有提升。如果要做INT8量化可以用ONNX Runtime的量化工具。量化需要校准数据集一般用几百张训练图像就够了。量化后的模型体积能减小到原来的1/4推理速度提升2到3倍但mIoU通常会下降1到3个百分点。如果对精度要求高可以用量化感知训练在训练时就模拟量化误差这样量化后的精度损失可以控制在0.5个百分点以内。5.3 ONNX模型的运行与验证ONNX模型可以用onnxruntime在Python里直接运行。创建InferenceSession时providers参数可以指定CUDAExecutionProvider或CPUExecutionProvider。如果要用TensorRT加速可以装onnxruntime-gpu然后用TensorRTExecutionProvider。运行时的输入预处理要和训练时一致包括归一化参数和图像尺寸。输出后处理就是argmax加颜色映射把类别索引转成可视化的分割图。我一般会写一个可视化脚本把原图、预测分割图、叠加图并排显示方便检查效果。6. 端侧部署TFLite、NCNN与RKNN的取舍6.1 TFLite转换与量化ONNX转TFLite需要经过TensorFlow。可以用onnx-tf工具先把ONNX转成TensorFlow SavedModel再用TFLiteConverter转成TFLite。这个过程比较容易出问题主要是算子兼容性。BiSeNet里的某些上采样和融合操作在TensorFlow里没有直接对应的算子需要手动替换。转换成功后可以用TFLite的量化工具做INT8量化。量化后的模型在移动端CPU上推理一张512x512的图像大概需要50到100毫秒具体取决于设备性能。如果设备支持GPU委托可以进一步加速。6.2 NCNN与RKNN的适用场景NCNN是腾讯开源的推理框架对ARM CPU优化很好。ONNX转NCNN可以用onnx2ncnn工具转换后需要手动修改param文件里的某些层参数。NCNN的优点是体积小、依赖少适合集成到Android或iOS应用里。RKNN是瑞芯微的推理框架专门针对RK系列芯片。如果你用的是RK3588这类开发板用RKNN可以获得最好的性能。ONNX转RKNN需要用RKNN-Toolkit2转换时要注意量化配置和目标平台设置。RKNN的量化精度通常比TFLite好一些因为它的量化工具更成熟。6.3 端侧推理的性能优化端侧推理的性能瓶颈通常在内存带宽和算子实现上。优化手段包括降低输入分辨率、减少模型通道数、使用深度可分离卷积、合并BN层到卷积层。我实测下来把输入从512x512降到256x256推理速度能提升3到4倍但mIoU会下降5到8个百分点。如果业务允许可以先用低分辨率做粗分割再在高分辨率上做局部精修。还有一个技巧把模型的输出从19类降到需要的类别。比如你只需要皮肤、头发、嘴唇三个区域可以在训练时就把其他类别合并这样输出通道从19降到3推理速度能提升不少。7. 常见问题与排查技巧实录7.1 训练不收敛或mIoU异常低最常见的原因是学习率太大或者数据标注有问题。先检查标注图里有没有全黑或者全白的异常样本再检查类别索引是否从0开始连续。如果标注没问题把学习率降到原来的1/10再试。另一个原因是损失函数的权重设置不合理。如果背景权重太高模型会倾向于预测背景。可以先把所有类别权重设成1.0跑几个epoch看看mIoU再逐步调整。7.2 ONNX导出失败或推理结果不对导出失败通常是算子不支持。可以尝试降低opset_version或者把模型里的某些操作替换成ONNX支持的等价操作。推理结果不对的话先检查输入预处理是否一致再检查输出是否需要softmax。BiSeNet的输出是logits如果训练时用了交叉熵损失推理时不需要额外softmax直接argmax就行。7.3 量化后精度下降严重量化精度下降的主要原因是校准数据集不够代表性。校准集应该覆盖各种光照、肤色、姿态的样本。如果下降超过5个百分点可以试试量化感知训练或者在量化时保留某些敏感层为FP16。7.4 端侧推理速度不达预期先确认是否用了正确的推理后端。TFLite在CPU上跑和用GPU委托跑速度差好几倍。NCNN要确认是否开启了Vulkan或OpenMP加速。RKNN要确认是否用了NPU而不是CPU。如果后端没问题检查模型是否有冗余算子。用Netron打开模型看看有没有可以合并的层。另外输入图像的预处理和后处理也可能成为瓶颈尽量把这些操作也放到GPU或NPU上。8. 一些实操心得与扩展思路我在实际项目里踩过最大的坑是数据标注的一致性。不同标注员对“头发”和“背景”的边界理解不一样导致模型学到的边界很模糊。后来我们写了一个标注规范文档明确了每个类别的边界定义还做了交叉验证mIoU直接提升了3个百分点。另一个心得是不要迷信预训练权重。CelebAMask-HQ上预训练的模型在你的业务数据上可能表现很差因为肤色、光照、拍摄设备都不一样。我通常会用预训练权重做初始化然后在自己的数据上fine-tune至少50个epoch。扩展思路方面BiSeNet的输出可以作为很多人脸编辑任务的基础。比如把嘴唇区域提取出来做颜色替换把头发区域提取出来做换发色把皮肤区域提取出来做磨皮。这些应用的核心都是精确的语义分割。如果你要做实时视频处理可以把BiSeNet和轻量级的人脸检测网络串联先检测再解析这样能保证每帧的处理速度。最后分享一个小技巧推理时可以用滑动窗口加多尺度融合。把图像缩放到不同尺度分别推理然后把结果平均能提升1到2个百分点的mIoU代价是推理时间翻倍。如果对精度要求高且算力充足值得一试。
返回列表