Ch24: ncnn 实战——编译、推理、性能测试

// classify.cpp
#include <stdio.h>
#include <algorithm>
#include <vector>
#include "net.h"

// 如果需要读取图片，可以用 stb_image（单头文件，零依赖）
#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"

struct ClassificationResult {
    int class_id;
    float score;
};

bool compare_score(const ClassificationResult& a, const ClassificationResult& b) {
    return a.score > b.score;
}

int classify(const char* model_param, const char* model_bin,
             const unsigned char* image_data, int img_w, int img_h) {
    // 1. 创建网络
    ncnn::Net net;
    // 可选：开启 Vulkan GPU 推理
    // net.opt.use_vulkan_compute = true;

    // 2. 加载模型
    if (net.load_param(model_param)) {
        fprintf(stderr, "load_param failed\n");
        return -1;
    }
    if (net.load_model(model_bin)) {
        fprintf(stderr, "load_model failed\n");
        return -1;
    }

    // 3. 准备输入
    // from_pixels_resize: 自动缩放到目标尺寸
    ncnn::Mat input = ncnn::Mat::from_pixels_resize(
        image_data,
        ncnn::Mat::PIXEL_RGB2BGR,  // 输入是 RGB，转为 BGR
        img_w, img_h,              // 原始图片尺寸
        224, 224                    // 目标尺寸
    );

    // 预处理参数（必须和训练时一致）
    const float mean_vals[3] = {104.f, 117.f, 123.f};
    const float norm_vals[3] = {1.f / 58.395f, 1.f / 57.12f, 1.f / 57.375f};
    input.substract_mean_normalize(mean_vals, norm_vals);

    // 4. 推理
    ncnn::Extractor ex = net.create_extractor();
    ex.set_num_threads(4);  // 设置推理线程数

    ex.input("input0", input);

    ncnn::Mat output;
    ex.extract("output0", output);

    // 5. 解析结果
    std::vector<ClassificationResult> results;
    for (int i = 0; i < output.w; i++) {
        ClassificationResult r;
        r.class_id = i;
        r.score = output[i];
        results.push_back(r);
    }

    // 按分数排序
    std::sort(results.begin(), results.end(), compare_score);

    // 输出 Top-5
    printf("Top-5 results:\n");
    for (int i = 0; i < 5 && i < (int)results.size(); i++) {
        printf("  class %d: %.4f\n", results[i].class_id, results[i].score);
    }

    return results[0].class_id;
}

int main(int argc, char** argv) {
    if (argc != 4) {
        fprintf(stderr, "Usage: %s <param> <bin> <image>\n", argv[0]);
        return -1;
    }

    // 读取图片
    int w, h, c;
    unsigned char* data = stbi_load(argv[3], &w, &h, &c, 3);
    if (!data) {
        fprintf(stderr, "Failed to load image: %s\n", argv[3]);
        return -1;
    }

    int result = classify(argv[1], argv[2], data, w, h);
    printf("Predicted class: %d\n", result);

    stbi_image_free(data);
    return 0;
}

ncnn::Net net;
net.load_param("model.param");
net.load_model("model.bin");

// 多线程场景：每个线程用自己的 Extractor
// Net 本身是只读的，可以安全地在多个线程间共享
#pragma omp parallel for
for (int i = 0; i < batch_size; i++) {
    ncnn::Extractor ex = net.create_extractor();
    ex.input("input0", inputs[i]);
    ncnn::Mat output;
    ex.extract("output0", output);
    // 处理 output...
}

// 假设模型有两个输出头：classification 和 detection
ncnn::Extractor ex = net.create_extractor();
ex.input("input0", input);

// 只需要 classification 结果——detection 分支不会被执行
ncnn::Mat cls_output;
ex.extract("classification_output", cls_output);

// 如果两个都需要：
ncnn::Mat det_output;
ex.extract("detection_output", det_output);
// 此时 detection 分支才会被计算（但 shared backbone 不会重复计算）

ncnn::Extractor ex = net.create_extractor();

// 开启轻量模式：推理完成后释放中间层的内存
// 代价是不能再次 extract 中间层的结果
ex.set_light_mode(true);

// 设置 blob 内存池分配器（减少内存碎片）
ncnn::PoolAllocator pool_allocator;
ncnn::UnlockedPoolAllocator unlocked_pool_allocator;
ex.set_blob_allocator(&pool_allocator);
ex.set_workspace_allocator(&unlocked_pool_allocator);

// bench_my_model.cpp
#include <stdio.h>
#include <chrono>
#include "net.h"

int main() {
    ncnn::Net net;
    net.load_param("my_model.ncnn.param");
    net.load_model("my_model.ncnn.bin");

    // 创建随机输入
    ncnn::Mat input(224, 224, 3);
    input.fill(0.5f);

    // 预热
    for (int i = 0; i < 3; i++) {
        ncnn::Extractor ex = net.create_extractor();
        ex.input("input0", input);
        ncnn::Mat output;
        ex.extract("output0", output);
    }

    // 正式测试
    int loop_count = 10;
    double total_ms = 0;

    for (int i = 0; i < loop_count; i++) {
        auto start = std::chrono::high_resolution_clock::now();

        ncnn::Extractor ex = net.create_extractor();
        ex.set_num_threads(4);
        ex.input("input0", input);
        ncnn::Mat output;
        ex.extract("output0", output);

        auto end = std::chrono::high_resolution_clock::now();
        double ms = std::chrono::duration<double, std::milli>(end - start).count();
        total_ms += ms;
        printf("Run %d: %.2f ms\n", i + 1, ms);
    }

    printf("Average: %.2f ms\n", total_ms / loop_count);
    return 0;
}