tensorrtx 实战：基于 TensorRT 网络定义 API 部署 ArcFace 人脸识别（R50 / R100 / MobileFaceNet）

本指南围绕 arcface/README.md 展开，系统讲解如何在 tensorrtx 仓库中把 InsightFace 生态下基于 MXNet 训练的 ArcFace 人脸识别模型（LResNet50E-IR、LResNet100E-IR 与 MobileFaceNet）转换为 TensorRT 可加载的权重文件、构建 engine 并完成推理，最终输出人脸 512/128 维特征向量与两张人脸图像的相似度得分。读完本文，你将掌握从 MXNet 权重到 TensorRT engine 的完整部署链路，以及项目中 BN 折叠、PReLU 实现等底层原理。

模块定位：tensorrtx 中的 arcface 目录

tensorrtx 是一个使用 TensorRT 网络定义 API（Network Definition API，而非 ONNX/TensorFlow 解析器）逐个复现经典深度学习网络的仓库。arcface 目录正是其中人脸识别（Face Recognition）场景的代表实现，其 MXNet 原始实现来自 deepinsight/insightface，README 明确说明本模块面向 TensorRT 8 环境。

该目录包含三类核心文件：

文件	作用
arcface-r50.cpp	LResNet50E-IR 网络的 engine 构建与推理入口
arcface-r100.cpp	LResNet100E-IR 网络的 engine 构建与推理入口
arcface-mobilefacenet.cpp	MobileFaceNet（轻量级）网络的 engine 构建与推理入口
gen_wts.py	把 MXNet 参数导出为 TensorRT 所需的 .wts 权重文件
prelu.cu / prelu.h	PReLU 自定义插件实现
CMakeLists.txt	三个可执行程序的构建脚本
logging.h / macros.h	TensorRT 日志与跨版本宏兼容

端到端部署流程总览

从 MXNet 预训练模型到可运行的人脸特征提取，整体分为三步：

1. 导出权重：把 MXNet 格式的预训练模型（model-r100-ii 等）通过 gen_wts.py 导出为 arcface-r100.wts（r50 / mobilefacenet 同理）；
2. 构建 engine：把 .wts 放到 tensorrtx/arcface 下，编译工程后用 ./arcface-r100 -s 序列化出 arcface-r100.engine；
3. 推理验证：用 ./arcface-r100 -d 反序列化 engine，对两张裁剪对齐后的人脸图像（joey0.ppm、joey1.ppm）做推理，输出各自耗时（latency）与相似度得分（similarity score）。

该流程的核心特点是全程使用 TensorRT 网络定义 API 逐层手工搭建网络（见 arcface-r100.cpp 的 createEngine 函数），不依赖任何模型解析器，因此对网络结构的每个算子都有完全可控的精度与性能设定。

第一步：从 MXNet 模型导出 .wts 权重文件

前置准备

• 克隆 deepinsight/insightface 仓库，README 建议 git checkout 3866cd77a6896c934b51ed39e9651b791d78bb57 这个提交；如果 master 分支异常，可回退到 94ad870abb3203d6f31b049b70dd080dc8f33fca；
• 下载对应的 MXNet 预训练模型压缩包（如 model-r100-ii.zip）并解压到 insightface/deploy 目录；
• 将本仓库的 gen_wts.py 复制到 insightface/deploy 目录，然后执行 python gen_wts.py。

gen_wts.py 原理

gen_wts.py 的逻辑非常简洁：加载 face_model.FaceModel 后，通过 model.model.get_params() 同时取出 MXNet 符号图的两类参数——args（权重，如卷积核、gamma 等）与 auxs（辅助量，如 BN 的 moving_mean / moving_var），并写入 .wts 文件：

f = open('arcface-r100.wts', 'w')
f.write('{}\n'.format(len(model.model.get_params()[0].keys()) + len(model.model.get_params()[1].keys())))
for k, v in model.model.get_params()[0].items():
    vr = v.reshape(-1).asnumpy()
    f.write('{} {} '.format(k, len(vr)))
    for vv in vr:
        f.write(' ')
        f.write(struct.pack('>f',float(vv)).hex())
    f.write('\n')

注意 struct.pack('>f', ...) 使用的是大端序（>）的 IEEE 754 单精度十六进制表示，这与 TensorRT C++ 端 loadWeights 的读取格式一一对应。该脚本默认参数 --model model-r100-ii/model,0，如需导出其他模型，可分别替换为对应 MXNet 模型前缀（r50 / mobilefacenet 同理类推）。

.wts 文件格式约定

TensorRT 端读取 .wts 的约定格式在 arcface-r100.cpp 的 loadWeights 中有清晰注释：

第一行：权重 blob 的总数量 count
后续每行：<name> <size> <以十六进制表示的 size 个 float 数据>

C++ 侧按此格式逐行解析，name 作为键存入 std::map<std::string, Weights>，size 为该张量的元素个数，后续数据用 std::hex 读取为 uint32_t 数组后按 float 解释。

第二步：编译工程并构建 engine

构建配置

CMakeLists.txt 使用 cmake_minimum_required(VERSION 2.6)、C++11 标准，并做如下关键配置：

• 通过 find_package(CUDA REQUIRED) 与 find_package(OpenCV) 链接 CUDA 与 OpenCV；
• 支持 aarch64 嵌入式平台：当 CMAKE_SYSTEM_PROCESSOR 匹配 aarch64 时自动切换 CUDA include/lib 路径到 /usr/local/cuda/targets/aarch64-linux；
• 将 prelu.cu 编译为共享库 myplugins，与 nvinfer、cudart、OpenCV_LIBS 一起链接到三个可执行程序 arcface-r50、arcface-r100、arcface-mobilefacenet。

编译与运行

cd tensorrtx/arcface
# 下载 joey0.ppm 和 joey1.ppm 放到 tensorrtx/arcface 目录
mkdir build
cd build
cmake ..
make
sudo ./arcface-r100 -s    # 序列化模型为 plan 文件 arcface-r100.engine
sudo ./arcface-r100 -d    # 反序列化 plan 文件并执行推理

r50 与 mobilefacenet 的可执行文件用法完全一致（-s 分别生成 arcface-r50.engine、arcface-mobilefacenet.engine）。需要注意：

• -s 与 -d 是 main 函数 中仅有的两个合法参数，其余参数会打印用法提示并以非零码退出；
• 工程默认以 maxBatchSize = 256 构建 engine（APIToModel(256, &modelStream)），但推理侧宏 BATCH_SIZE 固定为 1；
• .wts 与 .ppm 图像均放在 tensorrtx/arcface 目录，代码中通过 ../arcface-r100.wts、../joey0.ppm 这类相对路径访问（因为可执行文件位于 build/ 子目录）。

配置项详解

README 的 Config 小节指出两个宏开关，结合源码还可看到更多可调参数：

宏 / 常量	位置	默认值	说明
USE_FP16	三个 .cpp 文件顶部	注释状态（FP32）	取消注释即开启 FP16，构建时调用 config->setFlag(BuilderFlag::kFP16)
DEVICE	三个 .cpp 文件顶部	0	CUDA 设备 ID，main 中通过 cudaSetDevice(DEVICE) 生效
BATCH_SIZE	三个 .cpp 文件顶部	1	推理批大小，注释明确“currently, only support BATCH=1”
INPUT_H / INPUT_W	三个 .cpp 文件顶部	112 / 112	人脸输入尺寸，与 InsightFace 对齐后的裁剪图一致
OUTPUT_SIZE	r50 / r100 为 512，mobilefacenet 为 128	—	输出特征维度
INPUT_BLOB_NAME / OUTPUT_BLOB_NAME	三个 .cpp 文件顶部	"data" / "prob"	engine 的输入输出绑定名
最大 workspace	createEngine 末尾	16MB（16 * (1 << 20)）	通过 config->setMaxWorkspaceSize 设定

FP16 开关的实际生效位置在 createEngine 中：

#ifdef USE_FP16
    config->setFlag(BuilderFlag::kFP16);
#endif

源码级原理：网络是如何逐层搭建的

BatchNorm 折叠进 Scale 层

addBatchNorm2d 函数（arcface-r100.cpp）把 MXNet 的 BatchNorm 四个参数 gamma / beta / moving_mean / moving_var 按如下公式预折叠成 TensorRT IScaleLayer 的 scale、shift、power 三组权重：

scale = gamma / sqrt(var + eps)
shift = beta - mean * gamma / sqrt(var + eps)
power = 1.0

这样推理阶段 BN 不再单独占用算子，直接由一个通道级 Scale 完成，同时把 epsilon（网络各 BN 使用的 2e-5 或 MobileFaceNet 使用的 1e-3）作为显式参数传入。这是典型的“训练时 BN、推理时融合”的部署优化手法。

PReLU 的两种实现

InsightFace 网络大量使用 PReLU 激活，而 TensorRT 原生算子中 PReLU 支持受限，仓库给出两种方案：

1. Scale 组合方案（实际构建网络时使用）：addPRelu（arcface-r100.cpp）用 ReLU(x) + (-gamma) * ReLU(-x) 的数学等价关系，通过两次 ReLU、两次通道级 Scale 和一次 ElementWise 求和组合出 PReLU(x) = max(x,0) - gamma·max(-x,0)，全程无需自定义算子；
2. CUDA 插件方案：prelu.cu 实现了 PReluPlugin（IPluginV2IOExt），其核心 kernel prelu_kernel 逐元素计算 input >= 0 ? input : input * gamma[c]，配合 PluginCreator 支持从序列化流反序列化，被 CMake 编译为 myplugins 共享库供扩展使用。

ResNet IR 残差单元

R50 / R100 共用 resUnit 函数（arcface-r100.cpp）搭建 InsightFace 的 IR（Improved Residual）块：先 BN → 3×3 Conv → BN → PReLU → 3×3 Conv → BN，再做残差相加；当输入输出尺寸/通道不匹配（dim_match == false）时，用 1×1 卷积 + BN 的 shortcut 分支对齐后相加。R100 网络在 createEngine 中按 stage1(3) → stage2(13) → stage3(30) → stage4(3) 共 49 个残差单元堆叠，最终经过全局 BN 与 pre_fc1 全连接输出 512 维特征；R50 则是 stage1(3) → stage2(4) → stage3(14) → stage4(3)。

MobileFaceNet 的深度可分离结构

arcface-mobilefacenet.cpp 的网络则完全不同：DepthWise 用 setNbGroups(groups) 把卷积拆成 1×1 点卷积（pointwise）+ 3×3 深度卷积（depthwise）+ 1×1 投影卷积（projection）三段，DWResidual 再做残差求和；主干依次经过 conv_1、conv_2_dw、三段 res_x_block 残差块与 7×7 深度卷积，最终由 fc1 输出 128 维特征。这也是三个模型里唯一输出维度不同的原因。

engine 构建与推理调用链

• 序列化：APIToModel → createEngine → buildEngineWithConfig → engine->serialize()，写入 arcface-r100.engine；
• 反序列化：createInferRuntime → deserializeCudaEngine 读回 engine → 创建 IExecutionContext；
• 推理：doInference（arcface-r100.cpp）按绑定名拿到输入/输出 buffer，cudaMalloc 分配显存后通过 cudaMemcpyAsync + context.enqueue + cudaStreamSynchronize 完成一次异步推理，用时通过 std::chrono 计时输出。

输入预处理与相似度计算

人脸图像的预处理

推理输入是 112×112 的裁剪对齐人脸图（README 明确说明其由 insightface/deploy/face_model.py 的 get_input() 产生），测试图 joey0.ppm / joey1.ppm 需按 main 函数 中的方式处理：

cv::Mat img = cv::imread("../joey0.ppm");
for (int i = 0; i < INPUT_H * INPUT_W; i++) {
    data[i] = ((float)img.at<cv::Vec3b>(i)[2] - 127.5) * 0.0078125;        // R 通道
    data[i + INPUT_H * INPUT_W] = ((float)img.at<cv::Vec3b>(i)[1] - 127.5) * 0.0078125; // G 通道
    data[i + 2 * INPUT_H * INPUT_W] = ((float)img.at<cv::Vec3b>(i)[0] - 127.5) * 0.0078125; // B 通道
}

要点有二：一是 OpenCV 读入为 BGR 顺序，代码按通道重组回 RGB；二是归一化采用 InsightFace 标准的 (pixel - 127.5) * 0.0078125（即减 127.5 后除以 128），与训练侧预处理保持一致，避免精度偏移。

相似度得分

对 joey0 与 joey1 分别推理得到两个 512 维（mobilefacenet 为 128 维）特征向量后，代码用 cv::normalize 做 L2 归一化，再以矩阵点乘计算余弦相似度：

cv::Mat res = out_norm1 * out_norm;
std::cout << "similarity score: " << *(float*)res.data << std::endl;

程序运行后会依次打印两次推理各自的耗时（毫秒）与最终 similarity score，可据此直观验证部署正确性——同一人的两张对齐人脸应得到接近 1.0 的得分。

运行结果与验证要点

README 要求部署完成后检查三类输出：

1. engine 构建日志：构建时打印 Building engine, please wait for a while... 与 Build engine successfully!，若缺少 .wts 文件或格式错误会在 loadWeights 处断言失败；
2. 推理延迟（latency）：每次 doInference 前后计时，打印单次推理毫秒数；
3. 相似度得分（similarity score）：两张输入图的余弦相似度，用于验证模型功能正确性。

若结果明显异常，可优先检查：.wts 是否与当前可执行文件对应（r100 工程加载 arcface-r100.wts，不要混用）；输入图像是否经过与训练一致的对齐裁剪；预处理是否保持了 BGR→RGB 与 (x-127.5)*0.0078125 的约定。

常见问题与注意事项

• TensorRT 版本：本模块按 TensorRT 8 维护，macros.h 通过 NV_TENSORRT_MAJOR 区分 TRT_NOEXCEPT、TRT_CONST_ENQUEUE 等 API 差异，旧版本可能需自行适配；
• batch 限制：BATCH_SIZE 固定为 1，注释明确暂不支持更大 batch 的推理路径；
• FP16 与精度：默认 FP32 构建；开启 USE_FP16 可换取吞吐提升，但在精度敏感场景需结合相似度得分验证；
• 权重来源：预训练权重（ArcFace-R100 等）均来自 deepinsight/insightface 的 Model Zoo，本文不涉及具体下载链接，请以 README 中给出的来源为准；
• 嵌入式平台：在 aarch64 设备上编译时，CMake 会自动切换到 /usr/local/cuda/targets/aarch64-linux 的 CUDA 路径。

通过本文，你已经掌握了 tensorrtx/arcface 从 MXNet 权重导出、TensorRT 网络 API 搭建、engine 序列化/反序列化到人脸特征提取与相似度验证的完整闭环，可将同一套路迁移到 R50、R100 与 MobileFaceNet 三种常用人脸识别主干上。