MogFace嵌入式部署入门：从模型压缩到板端推理全流程

最近有不少朋友在问，训练好的AI模型怎么才能跑到像Jetson Nano或者RK3588这样的嵌入式板子上。确实，从云端服务器到巴掌大的设备，中间隔着好几道坎儿。今天，我就以人脸检测模型MogFace为例，带你走一遍完整的嵌入式部署流程。咱们不聊虚的，就讲怎么一步步把模型变小、变快，最终在板子上跑起来。整个过程有点像给模型“瘦身”和“搬家”，既要保证它还能认出人脸，又要适应新家的“小户型”和“低功耗”环境。

1. 环境准备与工具选择

在开始动手之前，得先把“工具箱”准备好。嵌入式部署和你在电脑上跑Python脚本完全是两码事，工具链的选择至关重要。

首先，你需要明确目标硬件平台。这决定了后续模型转换和代码编写的方向。目前主流的选择有两类：

• NVIDIA Jetson系列：比如Jetson Nano、Jetson Xavier NX。它们使用NVIDIA的GPU，推理引擎首选TensorRT。生态好，资料多，对新手相对友好。
• 其他AIoT芯片：比如瑞芯微的RK3588、RK3568，晶晨的A311D等。这些芯片通常有自家的推理SDK，比如瑞芯微的RKNN-Toolkit。

对于MogFace这个模型，我们假设它最初是用PyTorch训练好的一个.pth文件。我们的任务就是把这个文件，变成目标板上一个能高效运行的程序。

你需要准备的基本软件环境包括：

1. 模型训练环境：一台有GPU的Linux电脑（Ubuntu 18.04/20.04），用于最初的模型压缩和转换。这是我们的“加工车间”。
2. 交叉编译环境或板端开发环境：根据目标板选择。对于Jetson，可以在x86电脑上安装JetPack SDK进行交叉编译，也可以直接在板子上编译。对于RK3588，通常需要在x86电脑上配置RKNN的开发环境。
3. 目标板本身：准备好刷好官方系统镜像的板子，并通过SSH连接到它，方便我们上传文件和测试。

别被这些工具吓到，我们一步步来，每一步我都会给出具体的操作。

2. 第一步：模型压缩（瘦身计划）

直接从训练服务器上拿下来的模型，对于嵌入式设备来说通常都太“胖”了。动辄几百兆，内存吃不消，算力也跟不上。所以，部署前必须先“瘦身”。主要有两招：剪枝和蒸馏。

2.1 模型剪枝：去掉不重要的部分

你可以把神经网络想象成一棵茂密的大树。剪枝就是剪掉一些对最终结果影响不大的枝叶（神经元或连接），让树的结构更精简，但依然能开花结果。

对于MogFace这样的人脸检测模型，我们可以尝试对卷积层的通道进行剪枝。这里有一个非常简单的基于L1范数的通道剪枝示例，帮助你理解原理：

import torch
import torch.nn as nn

def channel_prune(model, prune_rate=0.3):
    """
    一个简单的通道剪枝函数示例。
    model: 要剪枝的模型
    prune_rate: 剪枝比例，例如0.3表示剪掉30%的通道
    """
    model.cpu()
    for name, module in model.named_modules():
        # 主要对卷积层进行剪枝
        if isinstance(module, nn.Conv2d):
            weight = module.weight.data # 形状: [out_channels, in_channels, k, k]
            # 计算每个输出通道的权重绝对值之和 (L1范数)
            channel_l1_norm = weight.abs().sum(dim=(1,2,3))
            # 确定要保留的通道索引
            num_keep = int(len(channel_l1_norm) * (1 - prune_rate))
            _, keep_indices = torch.topk(channel_l1_norm, num_keep)
            # 这里需要构建新的卷积层并复制权重，是一个简化说明
            print(f"Pruning layer {name}: {weight.size(0)} -> {num_keep} channels")
    # 注意：实际剪枝需要更复杂的处理，包括重建模型和调整后续层。
    return model

重要提示：上面的代码只是一个原理演示。实际工程中，你需要使用更成熟的剪枝库（如torch.nn.utils.prune或pytorch-model-compression），并且剪枝后必须进行微调，让模型重新适应，否则精度会掉得很厉害。你可以先尝试一个很小的剪枝比例（比如10%），微调后再评估精度损失。

2.2 知识蒸馏：让小模型学大模型

有时候，单纯剪枝效果有限。这时可以用知识蒸馏。它的思想是：让一个已经压缩好的“小模型”（学生），去学习原来那个庞大但精度高的“大模型”（老师）的输出行为，而不仅仅是学习原始数据标签。

对于MogFace，我们可以训练一个轻量级的Backbone（比如MobileNetV3）作为学生模型，让它的输出尽可能接近原始ResNet-based的MogFace（老师模型）的输出。这样，轻量级模型就能获得接近大模型的性能。

压缩完成后，你会得到一个更小的PyTorch模型文件。别忘了在测试集上验证一下，确保人脸检测的精度（mAP）还在可接受范围内。牺牲一点点精度，换来数倍的体积和速度提升，在嵌入式场景下是非常划算的买卖。

3. 第二步：模型格式转换（翻译成板子能懂的语言）

板子上的推理引擎（如TensorRT、RKNN）不认识PyTorch的.pth文件。我们需要把模型转换成它们认识的格式。这个步骤通常称为“模型转换”或“导出”。

3.1 通用中间格式：ONNX

在转换到最终格式前，我们常常先转到ONNX格式。它是一个开放的模型表示标准，相当于一个“通用翻译器”。

import torch
import onnx
from your_model_definition import MogFace # 假设这是你的模型定义类

# 加载压缩后的模型权重
model = MogFace()
model.load_state_dict(torch.load('pruned_mogface.pth'))
model.eval()

# 创建一个示例输入张量（模拟一张图片）
dummy_input = torch.randn(1, 3, 640, 640) # Batch=1, 3通道，高640，宽640

# 导出模型为ONNX格式
onnx_path = "mogface.onnx"
torch.onnx.export(
    model,
    dummy_input,
    onnx_path,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}} # 支持动态batch
)
print(f"Model exported to {onnx_path}")

导出ONNX后，建议用onnx.checker.check_model和onnxruntime进行验证，确保模型结构正确且能正常推理。

3.2 转换为目标引擎格式

有了ONNX这个“中转站”，我们就可以向最终目标进发了。

对于NVIDIA Jetson (TensorRT): 在你的JetPack环境（或配置了TensorRT的x86机器）上，使用trtexec工具或TensorRT的Python API进行转换。trtexec命令相对简单：

# 在Jetson板子上或装有TensorRT的机器上执行
trtexec --onnx=mogface.onnx \
        --saveEngine=mogface.plan \
        --workspace=1024 \
        --fp16 # 如果硬件支持FP16，可以显著加速

这条命令会将mogface.onnx转换为TensorRT引擎文件mogface.plan，并尝试使用FP16精度来提升速度。

对于RK3588 (RKNN): 你需要使用瑞芯微提供的RKNN-Toolkit2。这是一个Python工具包，通常在x86开发机上使用。

from rknn.api import RKNN

rknn = RKNN()

# 配置模型预处理、量化等参数
ret = rknn.config(mean_values=[[123.675, 116.28, 103.53]],
                  std_values=[[58.395, 57.12, 57.375]],
                  target_platform='rk3588')
# 加载ONNX模型
ret = rknn.load_onnx(model='mogface.onnx')
# 构建RKNN模型
ret = rknn.build(do_quantization=True, dataset='./dataset.txt') # 量化可以进一步压缩和加速
# 导出RKNN模型文件
ret = rknn.export_rknn('./mogface.rknn')

注意，量化（do_quantization=True）需要提供一个校准数据集（dataset.txt里是图片路径列表），用于统计激活值分布，是提升板端推理速度的关键一步。

4. 第三步：板端推理代码编写（让模型跑起来）

模型转换好了，接下来就要在板子上写程序调用它。这里我们分别给出TensorRT和RKNN的C++推理代码骨架。

4.1 Jetson平台TensorRT C++推理

在Jetson上，我们通常用C++来获得最佳性能。

// 示例代码骨架，展示关键步骤
#include <NvInfer.h>
#include <NvOnnxParser.h>
#include <iostream>
#include <fstream>
#include <vector>

class MogFaceTRT {
private:
    nvinfer1::IRuntime* runtime;
    nvinfer1::ICudaEngine* engine;
    nvinfer1::IExecutionContext* context;
    // ... 其他成员变量，如输入输出缓冲区指针

public:
    bool loadEngine(const std::string& enginePath) {
        std::ifstream file(enginePath, std::ios::binary);
        file.seekg(0, std::ios::end);
        size_t size = file.tellg();
        file.seekg(0, std::ios::beg);
        std::vector<char> engineData(size);
        file.read(engineData.data(), size);

        runtime = nvinfer1::createInferRuntime(logger); // logger需要实现
        engine = runtime->deserializeCudaEngine(engineData.data(), size);
        context = engine->createExecutionContext();
        // ... 分配输入输出CUDA内存
        return true;
    }

    void inference(const cv::Mat& inputImage) { // 假设使用OpenCV读图
        // 1. 图像预处理 (resize, normalize, HWC -> CHW, BGR -> RGB等)
        // 2. 将预处理后的数据从CPU内存拷贝到之前分配的GPU输入缓冲区
        // 3. 执行推理: context->executeV2(buffers); (或 enqueueV2)
        // 4. 将输出从GPU缓冲区拷贝回CPU内存
        // 5. 后处理: 解析输出张量，得到人脸框和关键点
        std::vector<FaceBox> faces = postprocess(outputData);
        // ... 绘制结果
    }

    ~MogFaceTRT() {
        // 按顺序释放资源: context, engine, runtime, CUDA内存等
    }
};

实际项目中，你需要处理繁琐的预处理/后处理、内存管理、错误处理等。可以借鉴NVIDIA官方示例代码。

4.2 RK3588平台RKNN C++推理

RKNN SDK也提供了C++接口。

#include <rknn_api.h>
#include <opencv2/opencv.hpp>
#include <iostream>

int main() {
    const char* model_path = "./mogface.rknn";
    rknn_context ctx;
    int ret;

    // 1. 加载RKNN模型
    ret = rknn_init(&ctx, model_path, 0, 0, nullptr);
    if (ret < 0) {
        std::cerr << "rknn_init failed: " << ret << std::endl;
        return -1;
    }

    // 2. 获取模型输入输出信息
    rknn_input_output_num io_num;
    ret = rknn_query(ctx, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
    // ... 获取具体的输入输出属性（维度、格式等）

    // 3. 准备输入数据
    cv::Mat img = cv::imread("test.jpg");
    cv::Mat resized, normalized;
    // ... 图像预处理 (resize, 归一化， BGR2RGB等)
    // 创建输入数据结构
    rknn_input inputs[1];
    inputs[0].index = 0;
    inputs[0].type = RKNN_TENSOR_UINT8; // 根据量化类型调整
    inputs[0].fmt = RKNN_TENSOR_NHWC;
    inputs[0].buf = normalized.data;
    inputs[0].size = normalized.total() * normalized.elemSize();
    ret = rknn_inputs_set(ctx, io_num.n_input, inputs);

    // 4. 执行推理
    ret = rknn_run(ctx, nullptr);

    // 5. 获取输出
    rknn_output outputs[io_num.n_output];
    // ... 为outputs分配内存
    ret = rknn_outputs_get(ctx, io_num.n_output, outputs, nullptr);

    // 6. 后处理
    // outputs[0].buf 里就是推理结果，需要根据模型定义解析为人脸框
    std::vector<FaceBox> faces = parse_output(outputs[0].buf, ...);
    // ... 绘制结果

    // 7. 释放资源
    rknn_outputs_release(ctx, io_num.n_output, outputs);
    rknn_destroy(ctx);
    return 0;
}

RKNN的C++ API调用流程相对直观，重点是处理好数据格式（量化后通常是UINT8）的匹配。

5. 第四步：性能调优与测试（精益求精）

模型跑起来不是终点，跑得好才是。嵌入式设备资源紧张，性能调优必不可少。

1. 调整推理线程数：对于有多核CPU的板子（如RK3588有4个A76大核和4个A55小核），可以通过设置推理引擎的线程数来充分利用CPU。在RKNN中，可以在rknn_init前通过rknn_set_core_mask API绑定核心。在TensorRT中，多线程调度更多依赖于你启动的推理线程本身。
2. 功耗模式选择：像Jetson Nano这样的板子，通常有几种功耗模式（5W, 10W Max-N, 10W Max等）。更高的功耗模式意味着更高的CPU/GPU频率和更好的性能，但发热也更大。你需要根据实际应用场景（持续运行还是间歇运行）和散热条件来选择。可以通过sudo nvpmodel -q查看和设置。
3. Pipeline优化：对于视频流处理，不要等一帧处理完再读下一帧。可以采用生产者-消费者流水线模式：一个线程负责抓取图像和预处理，另一个线程负责推理，第三个线程负责后处理和显示。这样可以最大化吞吐量。
4. 内存复用：频繁申请释放内存会产生开销。在初始化时就分配好所需的输入输出缓冲区，在整个程序生命周期内复用它们。
5. 量化验证：如果使用了INT8量化，务必在板端用充足的测试图片验证精度，确保量化没有引入不可接受的误差。

测试时，不仅要看单张图片的推理耗时，更要关注在模拟真实场景下的平均帧率、峰值内存占用以及长时间运行的稳定性。用tegrastats（Jetson）或top、free等命令监控板子状态。

---

获取更多AI镜像

想探索更多AI镜像和应用场景？访问 CSDN星图镜像广场，提供丰富的预置镜像，覆盖大模型推理、图像生成、视频生成、模型微调等多个领域，支持一键部署。