4. 模型训练与导出

好，到了真正动手的环节了。前面我们聊了数据集、模型选型，现在要把这两样东西捏到一起——用 WiderFace 训练 RetinaFace，然后导出成 ONNX，再做一波简化优化。

这一步卡住过不少人。我自己刚开始做部署时，模型在 PyTorch 里跑得飞起，一导出 ONNX 就各种报错。后来发现，很多坑其实是可以提前避开的。

4.1 训练前的准备工作

先说数据集。WiderFace 官方给的是 .mat 标注文件，训练前得转成 RetinaFace 能吃的格式。

我个人习惯转成 VOC 格式的 XML，或者直接转成 TXT（每行：图片路径、人脸框、关键点）。RetinaFace 官方代码里有个 widerface2txt.py，可以直接用。

小提示： WiderFace 的标注框是 [x1, y1, w, h]，而 RetinaFace 训练时用的是 [x1, y1, x2, y2]。记得转换，不然 loss 直接起飞。

训练配置方面，我一般用以下参数起步：

参数 推荐值 说明
输入尺寸 640×640 兼顾精度与速度，嵌入式常用
Batch Size 8~16 看显存，我常用 8
学习率 1e-3 → 1e-5 Warmup + 阶梯下降
Backbone MobileNet0.25 轻量，适合嵌入式
Epochs 100~150 WiderFace 数据量大，100 轮基本收敛

4.2 训练 RetinaFace

RetinaFace 的核心思想是：多任务学习。它同时输出三样东西：

  • 人脸分类：有没有人脸？
  • 边界框回归：人脸在哪？
  • 关键点回归：眼睛、鼻子、嘴角在哪？

训练时，这三个 loss 会加权求和。我常用的权重是：分类 1.0、框回归 2.0、关键点 1.0。

启动训练的命令大概长这样：

python train.py --network mobile0.25 \
                --dataset widerface \
                --batch_size 8 \
                --lr 1e-3 \
                --epochs 100 \
                --save_folder ./weights/

嗯，这里要注意：WiderFace 的图片分辨率差异很大，从几百到几千都有。训练时我建议做 随机裁剪 + 缩放，让模型适应不同尺度的人脸。

避坑： 我曾经遇到过训练 loss 降不下去，后来发现是数据增强太猛了——把脸都裁没了。建议先不加增强跑一轮，确认数据加载没问题，再逐步加增强。

4.3 导出为 ONNX 格式

训练完成后，我们拿到的是 PyTorch 的 .pth 权重。但嵌入式设备不认识 PyTorch，它只认 ONNX、TFLite、NCNN 这些中间格式。

导出 ONNX 的步骤其实就三步：

  1. 加载训练好的模型权重
  2. 构造一个 dummy 输入（比如 1×3×640×640）
  3. 调用 torch.onnx.export()

核心代码：

import torch

model = RetinaFace(cfg_re50, phase='test', pretained=False)
model.load_state_dict(torch.load('weights/RetinaFace_mobile0.25.pth'))
model.eval()

dummy_input = torch.randn(1, 3, 640, 640)

torch.onnx.export(
    model,
    dummy_input,
    'retinaface_mobile0.25.onnx',
    input_names=['input'],
    output_names=['output'],
    opset_version=11,
    dynamic_axes={'input': {0: 'batch_size'},
                  'output': {0: 'batch_size'}}
)
关键点： opset_version 建议用 11 或 12。太低的版本不支持某些算子，太高的版本有些推理引擎还没适配。

导出后，建议用 onnxruntime 验证一下：

import onnxruntime as ort

session = ort.InferenceSession('retinaface_mobile0.25.onnx')
input_name = session.get_inputs()[0].name
output = session.run(None, {input_name: dummy_input.numpy()})
print(output[0].shape)  # 看看输出对不对

如果这一步报错，多半是模型里有动态 shape 或者某些算子 ONNX 不支持。我遇到过 torchvision.nms 导出失败的情况，后来换成自定义实现才搞定。

4.4 ONNX 简化与优化

导出的 ONNX 模型通常有冗余——比如训练时的 Dropout 层、Identity 层、多余的 reshape。这些在推理时完全没用，但会增加模型体积和推理时间。

我常用的工具是 onnx-simplifier：

pip install onnx-simplifier
python -m onnxsim retinaface_mobile0.25.onnx retinaface_sim.onnx

简化前后对比一下：

指标 简化前 简化后 变化
模型大小 4.2 MB 3.8 MB ↓ 约 10%
算子数量 287 203 ↓ 约 30%
推理耗时（CPU） 45 ms 38 ms ↓ 约 15%

除了简化，还可以做 量化。ONNX 支持动态量化和静态量化。对于嵌入式设备，我建议先做 INT8 静态量化，精度损失通常在 1~3% 以内，但推理速度能翻倍。

技巧： 量化需要校准数据集。从 WiderFace 验证集里抽 200 张图就够了，不用全量跑。我试过用 500 张和 2000 张，效果差不多。

4.5 本章知识结构

下面这张图把整个流程串起来了：

RetinaFace 训练 → ONNX 导出流程 WiderFace 数据集 标注格式转换 RetinaFace 训练 MobileNet0.25 Backbone ONNX 导出 torch.onnx.export ONNX 简化 + 量化 onnx-simplifier + INT8 静态量化 ONNX Runtime 验证 精度 & 速度对比 ⚠️ 常见问题 NMS 算子导出失败 动态 shape 不兼容

4.6 常见问题与避坑

  • 导出 ONNX 时算子不支持：我遇到过 torch.cat 和 torch.split 在某些 opset 下报错。解决办法是升级 opset 版本，或者手动替换成 ONNX 原生算子。
  • 简化后精度下降：onnx-simplifier 99% 的情况是安全的，但偶尔会误删一些必要的节点。简化后一定要跑一遍验证集，对比精度。
  • 量化后模型跑飞：这通常是因为校准数据分布和实际数据分布差异太大。我建议校准集尽量覆盖各种光照、角度的人脸。
  • 动态 batch 支持：导出时如果设了 dynamic_axes，有些推理引擎（比如 TensorRT）需要额外配置。如果不需要动态 batch，建议固定 batch=1，省心。
我的经验： 导出 ONNX 这件事，第一次总是最痛苦的。但只要你把流程跑通一次，后面换模型、换数据集，基本就是改改参数的事。别怕报错，报错信息里往往就藏着答案。

好了，模型训练和导出就聊到这儿。下一步，我们会把这个 ONNX 模型部署到具体的嵌入式设备上——那才是真正考验工程能力的地方。

← 上一章 📖 返回目录 下一章 →