人脸解析face parsing这几年在美颜、AR试妆和人脸编辑项目里越来越常见本质上是把人脸按五官区域做一个像素级的语义分割。我最近把 BiSeNet 这套经典方案从 PyTorch 训练到 ONNX 导出、再到 TensorRT 部署完整跑通了一遍顺便把 19 类人脸语义分割的细节重新梳理了一轮。这篇文章把整个链路展开讲讲适合正在做人脸编辑、美颜、虚拟试妆这类项目的同学也适合想搞明白“语义分割模型到底怎么落到 GPU 上跑起来”的初学者。先说结论BiSeNet 在 512×512 输入下RTX 3090 上用 TensorRT FP16 推理能稳定跑到 1.5ms 左右而 PyTorch 原始推理大约是 6ms优化后提升明显。整个项目不复杂但里面坑不少尤其是数据预处理、类别索引重映射、ONNX 动态轴这几块稍微不注意就会让你排查半天。下面按我的实际操作顺序来写。1. 项目整体设计与技术选型思路1.1 19类人脸解析到底在解决什么问题人脸解析face parsing可以理解为人脸版的语义分割输入一张人脸图片模型输出一张同尺寸的分割图每个像素被标记为 19 个类别之一。这 19 类覆盖了五官、头发、衣服、帽子等关键部件具体分布如下索引类别说明索引类别说明0background背景10nose鼻子1skin皮肤11mouth嘴巴内部2l_brow左眉毛12u_lip上嘴唇3r_brow右眉毛13l_lip下嘴唇4l_eye左眼14neck脖子5r_eye右眼15neck_l衣领区域6eye_g眼镜16cloth衣服7l_ear左耳17hair头发8r_ear右耳18hat帽子9ear_r耳饰拿到这样的分割图之后上层应用就非常灵活了。比如美颜 App 要涂口红只需要把 u_lip 和 l_lip 两个类别的 mask 取出来做一个柔化上色换发色就锁定 hair 类背景替换可以直接把 background 抠掉。虚拟试妆、换脸、数字人贴图、人脸识别前的遮挡过滤几乎都能靠这份 19 类 mask 扩展出来。也就是说这个项目本身是一个通用的人脸理解前置模块而不是某个单一功能的实现。1.2 为什么选BiSeNet而不是U-Net或DeepLab我在最初选型时对比过 U-Net、DeepLabV3 和 BiSeNet。U-Net 虽然结构简单、训练稳定但编码器-解码器的结构对高分辨率输入比较吃力每一层特征都要做上采样和拼接计算量直线上升实时性不好。DeepLabV3 的空洞卷积ASPP模块确实能捕捉多尺度上下文但在显存占用和推理速度上都不算最优尤其是后面要部署到边缘设备时模型体积和延迟会让人头疼。BiSeNetBilateral Segmentation Network是专门为实时语义分割设计的。它的核心思想是“双边”一条空间路径保留高分辨率细节一条上下文路径提取全局语义信息最后通过特征融合模块合到一起。在人脸解析这个场景里细节非常重要——眉毛、嘴唇边缘稍微糊一点后面美颜效果就没法看。BiSeNet 的空间路径专门保细节所以很适合这种“既要快又要细”的任务。另外它的骨干网络用的是轻量化的 ResNet18/Xception模型体积小导出部署也方便。我当时还考虑过用预训练的 face-parsing 开源模型直接跑不走训练流程。但实际项目里往往需要自定义类别、换分辨率、调整后处理逻辑甚至要在自己的数据分布上微调所以还是自己走一遍训练和部署比较可控。这也是我推荐大家做的事先跑通训练链路再谈部署优化。2. 环境准备与数据集处理2.1 基础环境与依赖安装先交代一下我的实验环境Ubuntu 20.04Python 3.8PyTorch 1.12.1CUDA 11.6GPU 是 RTX 3090。其实 8GB 显存的显卡也能跑只是 batch size 要适当调小。依赖安装按下面来就行pip install torch1.12.1 torchvision0.13.1 pip install opencv-python pillow numpy tqdm albumentations这里提醒一句albumentations 不是必需品但它做数据增强时能自动同步处理图像和 mask避免自己写回调省很多事。如果你习惯用 OpenCV 手动写增强那也可以不装。我实际项目里还是用 albumentations 多一些后面会讲到原因。CUDA 版本和 PyTorch 版本的匹配问题值得注意。我当时先装了 PyTorch 2.0结果发现和机器上某个旧版 CUDA 工具链不兼容训练时一直报 undefined symbol。换回 1.12.1 就稳了。我的建议是不要盲目追新版本部署和生产环境优先选经过验证的组合。2.2 数据集准备与预处理流程人脸解析最常用的公开数据集是 CelebAMask-HQ它提供了 30000 张高清人脸图和对应的像素级标注 mask。每张 mask 是一个单通道 PNG像素值对应类别索引但原始标注里不同文件的类别索引并不完全统一所以第一步必须做重映射。我从实际踩坑的角度把流程拆成四步下载 CelebAMask-HQ 数据集确认图片和 mask 一一对应。写一个重映射脚本把所有 mask 的像素值统一到 0~18 的 19 类索引区间。按 9:1 划分训练集和验证集生成 train.txt 和 val.txt每行是“图片路径 空格 mask路径”。检查类别分布确认没有大面积缺失类别比如某些图片里没有帽子那 hat 类在所有像素里都不会出现这是正常的。重映射代码大致长这样import cv2 import numpy as np from pathlib import Path # 原始标注中的类别值需要映射到 0-18 连续索引 # 这里以 CelebAMask-HQ 官方 19 类顺序重新编号 mapping { 0: 0, # background 1: 1, # skin 2: 2, # l_brow # ... 省略完整映射表 } def remap_mask(mask_path, output_path): mask cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) remapped np.zeros_like(mask) for src, dst in mapping.items(): remapped[mask src] dst cv2.imwrite(output_path, remapped)重映射这一步为什么必须做因为如果某个类别的索引是 10但你训练时设置类别数是 19模型输出的是 19 通道计算交叉熵时索引 10 会落到第 10 类而不是真实的“鼻子”类。这种错位问题非常隐蔽我当时第一次训练时模型输出一片乱后面才发现是索引重映射漏了一行。数据增强方面我建议不要做太激进的变化。人脸任务跟通用分割不一样翻转、小角度旋转、轻微缩放和颜色抖动就够了。尤其不要随便做随机裁剪再拉伸因为这会破坏五官比例模型学到的几何特征会变歪。我用 albumentations 的组合是import albumentations as A train_transform A.Compose([ A.HorizontalFlip(p0.5), A.ShiftScaleRotate(shift_limit0.05, scale_limit0.1, rotate_limit10, p0.5), A.RandomBrightnessContrast(brightness_limit0.1, contrast_limit0.1, p0.3), A.Resize(512, 512) ])注意所有增强都要同步作用到 mask 上而且 mask 不能用灰度插值。像 ShiftScaleRotate 这类操作如果对 mask 用了默认的线性插值会产生介于两个类别之间的像素值训练时就会爆错或者静默错乱。albumentations 会自动把 mask 的插值方式设为最近邻这也是我推荐它的核心原因。3. BiSeNet核心结构解析3.1 空间路径与上下文路径是如何分工的BiSeNet 的结构可以用一个类比来理解你远远看一个人大脑先用整体轮廓判断“这是个穿红衣服的人”走近后又用眼睛捕捉“他眉毛旁边有一颗痣”。人脸解析也需要同时具备这两种能力——既要全图级的语义上下文也要保留像素级的轮廓细节。空间路径Spatial Path由三层卷积组成每一层 stride 都设为 1所以特征图分辨率不会降太低最终保持输入的 1/8 大小。这样做的好处是边缘、纹理信息损失很小眉毛、嘴唇这些精细结构能保住。缺点是感受野小单靠它不知道“这是左眼还是右眼”所以需要另一条路来补充语义信息。上下文路径Context Path用预训练的 ResNet18 或 Xception 作骨干网络。输入图片经过骨干网络逐层下采样得到一个语义很强的低分辨率特征再通过全局平均池化变成 1×1 的全局特征这相当于告诉模型“整张图我大概都看过了”。全局特征和中间层特征拼接之后会送入一个轻量的 ARM 模块。两条路径各自处理完之后空间路径输出的是 1/8 分辨率的细节特征上下文路径输出的是带有全局感受野的语义特征两者分辨率不同、语义层次不同不能直接相加这就轮到 FFM 模块来干活了。3.2 ARM模块与FFM融合的细节ARMAttention Refinement Module本质上是给上下文路径的特征加一个通道注意力。它先对输入特征做全局平均池化再用 1×1 卷积和 sigmoid 生成通道权重然后和原特征相乘最后加一个残差连接。这样做的目的是让模型自动决定哪些语义通道重要。比如在推理“帽子”时纹理特征通道可能不重要而颜色和边缘特征更重要权重就可以通过训练学出来。FFMFeature Fusion Module负责把空间路径和上下文路径的输出融合起来。典型实现是把两条路径的输出拼接然后分两路一路直接经过 1×1 卷积压缩通道另一路做全局池化加 1×1 卷积加 sigmoid 得到通道权重两路结果相乘得到融合特征。这相当于先降维再按通道重要性重新加权比简单 concat 加卷积高效得多。还有一个容易被忽略的设计辅助损失。BiSeNet 在上下文路径的两个中间阶段各接了一个分割头计算辅助损失然后和主损失加权相加。这样做可以让梯度更好地传到骨干网络底层加快收敛同时缓解深层网络训练时的梯度消失。我在训练时主损失权重设 1.0两个辅助损失权重各设 0.4效果比不加辅助损失稳定很多。我在理解整个结构时的一个心得是BiSeNet 之所以适合人脸解析不在于某一模块多精妙而在于它把“细节”和“语义”分离处理避免了这两个目标互相拖后腿。很多分割模型为了保细节会不断上采样结果语义信息稀释为了语义又会过度下采样结果边缘模糊。BiSeNet 让两条路径各干各的最后在 FFM 里统一这种思路在实际项目中非常实用。4. 模型训练实战与调参4.1 训练超参数与优化策略训练超参我基本沿用了 BiSeNet 原论文的设定再根据实际显卡做了小调整。配置如下输入分辨率512×512batch size8RTX 3090 上大约占 10GB 显存优化器SGDmomentum0.9weight_decay5e-4初始学习率0.01学习率策略poly即lr * (1 - iter / total_iter) ** 0.9迭代次数约 120k等价于 100 epochs 左右SGD 配合 poly 学习率在分割任务里一直很稳Adam 虽然收敛快但后期精度不如 SGD。人脸解析是密集预测在训练后期需要小学习率慢慢磨细节poly 策略比 StepLR 的阶梯式下降更适合。训练循环的核心代码我简化一下import torch import torch.nn as nn from torch.utils.data import DataLoader model BiSeNet(num_classes19, backboneresnet18) model.cuda() criterion nn.CrossEntropyLoss(ignore_index255) optimizer torch.optim.SGD(model.parameters(), lr0.01, momentum0.9, weight_decay5e-4) def poly_lr(epoch, total_epoch): return 0.01 * (1 - epoch / total_epoch) ** 0.9 for epoch in range(total_epoch): model.train() lr poly_lr(epoch, total_epoch) for param_group in optimizer.param_groups: param_group[lr] lr for images, labels in train_loader: images, labels images.cuda(), labels.cuda() outputs model(images) loss criterion(outputs[0], labels) # 主损失 loss 0.4 * criterion(outputs[1], labels) # 辅助损失 loss 0.4 * criterion(outputs[2], labels) optimizer.zero_grad() loss.backward() optimizer.step()注意 labels 的 shape 必须是[B, H, W]通道维不要留着。如果传入[B, 1, H, W]的 maskCrossEntropyLoss 会报维度错误。这是我第一次写训练脚本时踩过的坑印象特别深。4.2 训练过程与效果评估我在 CelebAMask-HQ 上训练时前 10 个 epoch loss 从 2.3 左右降到 0.8 左右速度很快但到 50 epoch 之后 loss 下降就非常缓慢了主要是在磨眉毛、嘴唇这些边界区域。最终验证集 mIoU 大约 0.90每类 IoU 取平均单类精度上 skin、background、hair 这些大区域很容易到 0.95 以上但 earrings、hat 这类出现频率低的类别 IoU 只有 0.7 左右属于正常现象。训练时还要关注一个点显存溢出。如果你用 RTX 2080Ti 这种 11GB 显卡batch size 8 再加辅助损失的反向传播显存可能会爆。我的建议是优先开 PyTorch AMP自动混合精度from torch.cuda.amp import autocast, GradScaler scaler GradScaler() with autocast(): outputs model(images) loss criterion(outputs[0], labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()AMP 在 3090 上能省 30%~40% 显存且对最终精度影响很小属于“加了不亏”的选项。还有一个小技巧是降低输入的 batch size 同时把分辨率保持 512因为人脸解析对分辨率比 batch size 更敏感。5. 模型部署与推理优化5.1 从PyTorch到ONNX的导出训练完成后部署阶段第一步是把 PyTorch 模型导出为 ONNX。这一步看起来简单但有几个细节不处理会折腾半天。首先必须把模型切成 eval 模式并禁用梯度因为 BN 层和 dropout 在推理时行为不同。其次ONNX 导出时要把输入输出维度写清楚。我的输入是[1, 3, 512, 512]输出是[1, 19, 512, 512]19 对应 19 个类别通道。最后因为实际应用里输入尺寸可能不固定需要设置 dynamic_axes。导出代码import torch model BiSeNet(num_classes19, backboneresnet18) checkpoint torch.load(bisenet_celeba.pth, map_locationcpu) model.load_state_dict(checkpoint[state_dict], strictFalse) model.eval() dummy_input torch.randn(1, 3, 512, 512) torch.onnx.export( model, dummy_input, bisenet_face.onnx, input_names[input], output_names[output], opset_version11, dynamic_axes{ input: {0: batch, 2: height, 3: width}, output: {0: batch, 2: height, 3: width} } )导出之后一定要用 onnxruntime 做一次精度比对import onnxruntime as ort import numpy as np sess ort.InferenceSession(bisenet_face.onnx, providers[CUDAExecutionProvider]) x np.random.randn(1, 3, 512, 512).astype(np.float32) onnx_out sess.run(None, {input: x})[0] with torch.no_grad(): torch_out model(torch.from_numpy(x)).cpu().numpy() print(max diff:, np.abs(onnx_out - torch_out).max())如果 max diff 在 1e-4 量级基本说明导出没问题。如果差得很多先检查模型里有没有自定义 op 或者训练模式的残留操作。我当时遇到过一次导出后输出全是 NaN排查半天发现是模型里有一个 BatchNorm 的 running_mean 没有被正确 load 进 checkpoint导出时 BN 层的统计量是随机的导致推理结果完全不可用。5.2 TensorRT加速与INT8量化实战ONNX 是中间格式真正部署到 GPU 上加速还需要转成 TensorRT。我用的是 trtexec 工具trtexec --onnxbisenet_face.onnx \ --saveEnginebisenet_fp16.engine \ --fp16 \ --minShapesinput:1x3x512x512 \ --optShapesinput:8x3x512x512 \ --maxShapesinput:16x3x512x512这里的 minShapes、optShapes、maxShapes 要跟 dynamic_axes 对应。FP16 精度下RTX 3090 实测单张 512×512 输入大约 1.5ms比 ONNX Runtime 的 4ms 快了不少比 PyTorch 原始推理 6ms 更是提升明显。如果还想再进一步压缩可以尝试 INT8 量化。TensorRT 的 INT8 需要校准数据集我用验证集里随机抽 500 张图作为校准数据生成校准缓存trtexec --onnxbisenet_face.onnx \ --saveEnginebisenet_int8.engine \ --int8 \ --calibratorhistogram \ --calibDatacalibration_dataINT8 量化后推理延迟能降到约 1ms但 mIoU 会从 FP16 的 0.90 掉到 0.87 左右主要损失集中在眉毛和眼镜这类细长区域。实际项目中要看场景是否允许这个精度损失。我个人建议优先用 FP16因为精度更稳部署排错成本低。TensorRT 的 Python 推理代码也不复杂import tensorrt as trt import pycuda.driver as cuda import pycuda.autoinit logger trt.Logger(trt.Logger.WARNING) with open(bisenet_fp16.engine, rb) as f: engine_data f.read() runtime trt.Runtime(logger) engine runtime.deserialize_cuda_engine(engine_data) context engine.create_execution_context() # 输入输出缓冲区绑定 input_buf cuda.mem_alloc(1 * 3 * 512 * 512 * 4) output_buf cuda.mem_alloc(1 * 19 * 512 * 512 * 4)推理循环里只需要把输入图像数据复制到 input_buf执行 context.execute_v2 后从 output_buf 拷回分割结果即可。注意输入图像必须是 RGB 顺序且一定要做和训练时一致的归一化操作否则输出会偏。6. 常见问题与排查技巧实录我把这一路踩过的坑整理成了速查表方便直接对照现象可能原因解决思路训练时显存 OOMbatch size 太大、未启用 AMP调小 batch size、开启混合精度、临时降分辨率到 448分割图整片是背景类类别索引重映射出错、label 增强时被插值破坏重映射表单独写个测试用例增强中 mask 用最近邻ONNX 导出时报维度错误未设置 dynamic_axes、模型里有固定尺寸操作设置动态轴、用 opset 11 以上必要时用 torch.onnx.export 的 verbose 查看节点ONNX 与 PyTorch 推理结果差异大BN 层统计量没加载、预处理不一致确认 checkpoint 完整加载统一归一化参数和 resize 插值TensorRT FP16 结果局部断裂FP16 精度不足细长区域受影响改用 FP32 引擎或对敏感类别做后处理修复推理速度上不去输入分辨率过大、batch 太小、CPU 解码是瓶颈降到 512 或 384、增大 batch、并行解码输出 mask 有黑边resize 时未处理 align_corners训练和部署统一 align_corners 设置除了表格里的问题还有几个经验值得单独说一下。第一预处理一致性怎么强调都不过分。训练时如果用 BGR 输入部署时就不能用 RGB训练时归一化是mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]部署时就必须一致。很多时候模型部署后效果不对不是模型坏了而是图片喂进去之前已经变了样。我为这件事专门写了一个preprocess.py训练和部署共用同一套函数从源头避免不一致。第二后处理不要小看。模型输出的 19 通道概率图需要一个 argmax 操作转成单通道索引再通过一个 19 色的 color map 转成可视化 PNG 才能看效果。很多人直接在 TensorRT 输出上 argmax却忘了 softmax 不是必须的——argmax 本身不受单调变换影响所以不需要先 softmax。这能省一部分计算。第三如果要在生产环境里处理视频流建议把 TensorRT 推理放到单独的进程或线程里输入用队列缓冲。因为 GPU 推理本身很快但图像解码、缩放、颜色转换这些 CPU 操作很容易成为瓶颈。我试过用 OpenCV 直接读视频逐帧推理FPS 只有 20 左右改成多线程解码加队列之后同样的模型能跑到 40 FPS 以上。最后再分享一个小技巧。如果你只是想做快速验证可以先下载别人训练好的 BiSeNet 人脸解析权重把 ONNX 导出和 TensorRT 部署链路跑通确认整个推理管线没有问题再回过头来用自己的数据集训练。这样能把“模型训练”和“工程部署”两件事解耦排查问题时心里更有底。我个人在实际操作中的体会是人脸解析这一类项目的难点其实不在模型结构而在数据与工程细节。类别索引错一位、预处理差一个插值方式、部署时动态轴多设少设每一个小问题都会让你排查半天。但只要把这套流程走顺了后续不管换模型、换数据集还是换部署平台都是一个套路。希望这篇实战记录能帮你少走几步弯路。
