如何使用 NCNN 框架在移动端高效部署 PyTorch 模型

在端侧 AI 落地过程中，如何让原本在服务器跑的重量级模型在手机端「跑得快、不发烫」是核心挑战。腾讯开发的 NCNN 是一个针对移动端优化的极致高性能神经网络推理框架，它无第三方依赖、针对 ARM 架构进行了深度的汇编指令优化。本文将带你实操如何将 PyTorch 模型部署到 NCNN。

1. 环境准备

首先，确保安装了 torch、onnx 和 onnx-simplifier。后者非常关键，能去掉 ONNX 中冗余的胶水算子。

pip install torch torchvision onnx onnx-simplifier

2. 导出并优化 ONNX 模型

以 ResNet18 为例，第一步是将其导出为标准 ONNX 格式，并进行简化。

import torch
import torchvision.models as models
from onnxsim import simplify

# 加载预训练模型
model = models.resnet18(pretrained=True).eval()
dummy_input = torch.randn(1, 3, 224, 224)

# 导出 ONNX
torch.onnx.export(model, dummy_input, "resnet18.onnx", 
                  input_names=["input"], output_names=["output"], 
                  opset_version=11)

# 使用 onnx-simplifier 简化模型，这一步对于 NCNN 的算子匹配至关重要
import onnx
onnx_model = onnx.load("resnet18.onnx")
model_simp, check = simplify(onnx_model)
if check:
    onnx.save(model_simp, "resnet18_sim.onnx")
    print("模型简化成功！")

3. 转换为 NCNN 格式

下载并编译好 NCNN 后，使用 onnx2ncnn 工具生成 NCNN 所需的 .param（模型结构）和 .bin（权重数据）文件。

./onnx2ncnn resnet18_sim.onnx resnet18.param resnet18.bin

4. 在端侧进行 C++ 推理测试

NCNN 的 API 简洁高效，以下是加载模型并进行推理的核心代码示例：

#include "net.h"
#include <algorithm>
#include <vector>

int main() {
    ncnn::Net resnet18;
    // 加载模型结构和权重
    resnet18.load_param("resnet18.param");
    resnet18.load_model("resnet18.bin");

    // 构造输入数据，注意 NCNN 使用自己的 Mat 格式
    ncnn::Mat in = ncnn::Mat::from_pixels_resize(image_data, ncnn::Mat::PIXEL_RGB, width, height, 224, 224);

    // 数据预处理：减均值、除方差 (Imagenet 标准)
    const float mean_vals[3] = {103.939f, 116.779f, 123.68f};
    const float norm_vals[3] = {0.017f, 0.017f, 0.017f};
    in.substract_mean_normalize(mean_vals, norm_vals);

    // 执行推理
    ncnn::Extractor ex = resnet18.create_extractor();
    ex.input("input", in);
    ncnn::Mat out;
    ex.extract("output", out);

    // 获取结果
    std::vector<float> scores;
    scores.resize(out.w);
    for (int j = 0; j < out.w; j++) {
        scores[j] = out[j];
    }
    return 0;
}

总结与建议

1. 算子对齐：若转换失败，检查是否使用了 NCNN 不支持的算子，尽量使用静态 Shape 导出。
2. 量化加速：对于追求极致速度的场景，可以使用 NCNN 自带的 ncnn2table 进行 INT8 量化。
3. 多线程：NCNN 默认开启多线程，可以通过 ncnn::set_cpu_powersave 设置核心运行策略。