实战案例7：机器学习推理插件（ONNX/TensorRT动态加载）

说实话，把机器学习模型塞进C语言程序里，这事儿我干过不少次。最早是在一个工业视觉检测项目里，甲方要求用TensorRT做推理加速，但整个系统是C写的嵌入式框架。你想想看，总不能为了跑个模型就把整个架构推倒重来吧？

所以动态加载就成了最自然的选择。模型推理作为一个插件，按需加载，按需卸载。今天我就拿ONNX Runtime和TensorRT这两个主流推理引擎，给你演示怎么把它们做成动态库插件。

为什么需要动态加载推理引擎

我遇到过不少团队，一开始把推理引擎静态链接进主程序。结果呢？每次升级模型框架都得重新编译整个项目。更麻烦的是，ONNX和TensorRT的依赖库经常打架——一个要libprotobuf 3.11，另一个非要3.20。

动态加载的好处很明显：

• 解耦：主程序不需要链接任何推理库
• 灵活：运行时切换推理后端，甚至热更新模型
• 隔离：每个插件有自己的依赖环境，互不干扰

统一的推理接口定义

先定义插件需要暴露的接口。我个人习惯用纯C的结构体函数指针表，这样ABI最稳定，C++的name mangling容易出问题。

// inference_plugin.h
#ifndef INFERENCE_PLUGIN_H
#define INFERENCE_PLUGIN_H

#include <stdint.h>

#ifdef __cplusplus
extern "C" {
#endif

// 推理引擎类型
typedef enum {
    INFERENCE_ENGINE_ONNX = 0,
    INFERENCE_ENGINE_TENSORRT,
    INFERENCE_ENGINE_UNKNOWN
} inference_engine_t;

// 张量数据结构
typedef struct {
    float* data;
    int32_t* shape;
    int32_t ndim;
    int32_t size;  // 元素总数
} tensor_t;

// 模型配置
typedef struct {
    const char* model_path;
    inference_engine_t engine_type;
    int32_t device_id;       // GPU设备ID
    int32_t batch_size;
    int32_t use_fp16;        // 是否使用半精度
} model_config_t;

// 插件接口结构体
typedef struct {
    // 插件信息
    const char* (*get_engine_name)(void);
    const char* (*get_version)(void);
    
    // 生命周期
    int32_t (*init)(const model_config_t* config);
    int32_t (*destroy)(void);
    
    // 推理
    int32_t (*infer)(tensor_t* input, tensor_t* output);
    
    // 查询
    int32_t (*get_input_shape)(int32_t* shape, int32_t* ndim);
    int32_t (*get_output_shape)(int32_t* shape, int32_t* ndim);
} inference_plugin_t;

// 插件必须导出的符号
extern inference_plugin_t plugin;

#ifdef __cplusplus
}
#endif

#endif

这里有个细节：plugin是一个全局变量，不是函数。为什么？因为dlsym获取符号地址时，全局变量比函数更直观，而且可以包含多个函数指针。我在项目里吃过亏——用函数返回结构体，结果不同编译器对返回值的处理方式不一样，踩了不少坑。

ONNX Runtime插件实现

ONNX Runtime的C API其实挺清爽的。不过要注意，它要求先创建环境（OrtEnv），再创建会话（OrtSession）。我习惯在init里一次性搞定。

// onnx_plugin.c
#include "inference_plugin.h"
#include <onnxruntime_c_api.h>
#include <stdlib.h>
#include <string.h>

static const OrtApi* g_ort = NULL;
static OrtEnv* g_env = NULL;
static OrtSession* g_session = NULL;
static OrtMemoryInfo* g_memory_info = NULL;

static const char* get_engine_name(void) {
    return "ONNX Runtime";
}

static const char* get_version(void) {
    return "1.15.0";
}

static int32_t init(const model_config_t* config) {
    // 获取ONNX Runtime API
    g_ort = OrtGetApiBase()->GetApi(ORT_API_VERSION);
    if (!g_ort) return -1;
    
    // 创建环境
    OrtStatus* status = g_ort->CreateEnv(ORT_LOGGING_LEVEL_WARNING, 
                                          "onnx_plugin", &g_env);
    if (status) {
        g_ort->ReleaseStatus(status);
        return -1;
    }
    
    // 创建会话选项
    OrtSessionOptions* session_options;
    g_ort->CreateSessionOptions(&session_options);
    g_ort->SetSessionGraphOptimizationLevel(session_options, 
                                             ORT_ENABLE_ALL);
    
    // 加载模型
    status = g_ort->CreateSession(g_env, config->model_path,
                                  session_options, &g_session);
    g_ort->ReleaseSessionOptions(session_options);
    
    if (status) {
        g_ort->ReleaseStatus(status);
        return -1;
    }
    
    // 创建内存信息
    g_ort->CreateMemoryInfo("Cpu", OrtDeviceAllocator, 0,
                            OrtMemTypeDefault, &g_memory_info);
    
    return 0;
}

static int32_t infer(tensor_t* input, tensor_t* output) {
    if (!g_session || !input || !output) return -1;
    
    // 创建输入张量
    OrtValue* input_tensor = NULL;
    OrtStatus* status = g_ort->CreateTensorWithDataAsOrtValue(
        g_memory_info, input->data, input->size * sizeof(float),
        input->shape, input->ndim, ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT,
        &input_tensor);
    if (status) return -1;
    
    // 运行推理
    OrtValue* output_tensor = NULL;
    const char* input_names[] = {"input"};
    const char* output_names[] = {"output"};
    
    status = g_ort->Run(g_session, NULL, input_names,
                         &input_tensor, 1, output_names, 1,
                         &output_tensor);
    g_ort->ReleaseValue(input_tensor);
    if (status) return -1;
    
    // 获取输出数据
    float* output_data;
    g_ort->GetTensorMutableData(output_tensor, (void**)&output_data);
    
    // 获取输出形状
    OrtTensorTypeAndShapeInfo* shape_info;
    g_ort->GetTensorTypeAndShape(output_tensor, &shape_info);
    g_ort->GetDimensions(shape_info, output->shape, &output->ndim);
    
    // 计算元素总数
    output->size = 1;
    for (int i = 0; i < output->ndim; i++) {
        output->size *= output->shape[i];
    }
    
    // 拷贝数据
    memcpy(output->data, output_data, output->size * sizeof(float));
    
    g_ort->ReleaseTensorTypeAndShapeInfo(shape_info);
    g_ort->ReleaseValue(output_tensor);
    
    return 0;
}

// ... 其他函数实现省略

inference_plugin_t plugin = {
    .get_engine_name = get_engine_name,
    .get_version = get_version,
    .init = init,
    .destroy = destroy,
    .infer = infer,
    .get_input_shape = get_input_shape,
    .get_output_shape = get_output_shape
};

TensorRT插件实现

TensorRT的API比ONNX Runtime复杂一些，因为它需要先构建引擎（engine），再创建执行上下文（execution context）。构建引擎这一步很慢，我一般会在init里从序列化文件加载，避免每次启动都重新构建。

// tensorrt_plugin.c
#include "inference_plugin.h"
#include <NvInfer.h>
#include <cuda_runtime_api.h>
#include <stdlib.h>
#include <string.h>

using namespace nvinfer1;

class Logger : public ILogger {
    void log(Severity severity, const char* msg) override {
        // 只打印警告和错误
        if (severity <= Severity::kWARNING) {
            fprintf(stderr, "[TensorRT] %s\n", msg);
        }
    }
};

static Logger g_logger;
static IRuntime* g_runtime = NULL;
static ICudaEngine* g_engine = NULL;
static IExecutionContext* g_context = NULL;
static cudaStream_t g_stream = NULL;
static void* g_buffers[2];  // 输入输出缓冲区

static int32_t init(const model_config_t* config) {
    // 创建runtime
    g_runtime = createInferRuntime(g_logger);
    if (!g_runtime) return -1;
    
    // 读取序列化引擎文件
    FILE* fp = fopen(config->model_path, "rb");
    if (!fp) return -1;
    
    fseek(fp, 0, SEEK_END);
    size_t size = ftell(fp);
    rewind(fp);
    
    char* engine_data = (char*)malloc(size);
    fread(engine_data, 1, size, fp);
    fclose(fp);
    
    // 反序列化引擎
    g_engine = g_runtime->deserializeCudaEngine(engine_data, size);
    free(engine_data);
    
    if (!g_engine) return -1;
    
    // 创建执行上下文
    g_context = g_engine->createExecutionContext();
    if (!g_context) return -1;
    
    // 创建CUDA流
    cudaStreamCreate(&g_stream);
    
    // 分配GPU内存
    int input_idx = g_engine->getBindingIndex("input");
    int output_idx = g_engine->getBindingIndex("output");
    
    Dims input_dims = g_engine->getBindingDimensions(input_idx);
    Dims output_dims = g_engine->getBindingDimensions(output_idx);
    
    size_t input_size = 1, output_size = 1;
    for (int i = 0; i < input_dims.nbDims; i++) input_size *= input_dims.d[i];
    for (int i = 0; i < output_dims.nbDims; i++) output_size *= output_dims.d[i];
    
    cudaMalloc(&g_buffers[0], input_size * sizeof(float));
    cudaMalloc(&g_buffers[1], output_size * sizeof(float));
    
    return 0;
}

static int32_t infer(tensor_t* input, tensor_t* output) {
    if (!g_context || !input || !output) return -1;
    
    // 拷贝输入数据到GPU
    cudaMemcpyAsync(g_buffers[0], input->data,
                    input->size * sizeof(float),
                    cudaMemcpyHostToDevice, g_stream);
    
    // 执行推理
    g_context->enqueueV2(g_buffers, g_stream, nullptr);
    
    // 拷贝输出数据回CPU
    cudaMemcpyAsync(output->data, g_buffers[1],
                    output->size * sizeof(float),
                    cudaMemcpyDeviceToHost, g_stream);
    
    cudaStreamSynchronize(g_stream);
    
    return 0;
}

// ... 其他函数实现省略

extern "C" {
    inference_plugin_t plugin = {
        .get_engine_name = get_engine_name,
        .get_version = get_version,
        .init = init,
        .destroy = destroy,
        .infer = infer,
        .get_input_shape = get_input_shape,
        .get_output_shape = get_output_shape
    };
}

主程序动态加载

主程序这边就简单了。用dlopen加载插件，dlsym获取plugin符号，然后调用接口就行。

// main.c
#include <dlfcn.h>
#include <stdio.h>
#include "inference_plugin.h"

int main(int argc, char* argv[]) {
    if (argc < 3) {
        fprintf(stderr, "Usage: %s <plugin.so> <model_path>\n", argv[0]);
        return 1;
    }
    
    // 加载插件
    void* handle = dlopen(argv[1], RTLD_LAZY | RTLD_LOCAL);
    if (!handle) {
        fprintf(stderr, "Failed to load plugin: %s\n", dlerror());
        return 1;
    }
    
    // 获取插件接口
    inference_plugin_t* plugin = dlsym(handle, "plugin");
    if (!plugin) {
        fprintf(stderr, "Failed to find plugin symbol\n");
        dlclose(handle);
        return 1;
    }
    
    printf("Loaded engine: %s v%s\n",
           plugin->get_engine_name(),
           plugin->get_version());
    
    // 配置模型
    model_config_t config = {
        .model_path = argv[2],
        .engine_type = INFERENCE_ENGINE_UNKNOWN,
        .device_id = 0,
        .batch_size = 1,
        .use_fp16 = 0
    };
    
    // 初始化
    if (plugin->init(&config) != 0) {
        fprintf(stderr, "Failed to init plugin\n");
        dlclose(handle);
        return 1;
    }
    
    // 准备输入数据
    tensor_t input = {
        .data = (float*)malloc(3 * 224 * 224 * sizeof(float)),
        .shape = (int32_t[]){1, 3, 224, 224},
        .ndim = 4,
        .size = 3 * 224 * 224
    };
    
    tensor_t output = {
        .data = (float*)malloc(1000 * sizeof(float)),
        .shape = (int32_t[]){1, 1000},
        .ndim = 2,
        .size = 1000
    };
    
    // 执行推理
    if (plugin->infer(&input, &output) == 0) {
        printf("Inference succeeded!\n");
        // 处理输出...
    }
    
    // 清理
    plugin->destroy();
    free(input.data);
    free(output.data);
    dlclose(handle);
    
    return 0;
}

编译与测试

编译插件时，记得用-fPIC和-shared。主程序链接-ldl。

# 编译ONNX插件
gcc -fPIC -shared -o onnx_plugin.so onnx_plugin.c \
    -I/path/to/onnxruntime/include \
    -L/path/to/onnxruntime/lib -lonnxruntime

# 编译TensorRT插件
g++ -fPIC -shared -o tensorrt_plugin.so tensorrt_plugin.c \
    -I/path/to/tensorrt/include \
    -L/path/to/tensorrt/lib -lnvinfer -lcudart

# 编译主程序
gcc -o inference_app main.c -ldl

# 测试
./inference_app ./onnx_plugin.so ./model.onnx
./inference_app ./tensorrt_plugin.so ./model.engine

插件加载架构图

下面这张图展示了整个动态加载推理插件的流程。主程序通过统一的接口与插件交互，底层可以是ONNX Runtime、TensorRT，或者将来任何新的推理引擎。

这张图其实反映了我做插件化架构的一个核心理念：接口稳定，实现灵活。主程序只跟inference_plugin_t打交道，至于背后是ONNX还是TensorRT，它根本不关心。哪天想换成OpenVINO或者TFLite，再写个新插件就行，主程序一行代码都不用改。

嗯，最后说一句。动态加载推理引擎这件事，看起来技术含量不高，但实际落地时坑不少。比如CUDA上下文的管理、多线程下的线程安全、GPU内存泄漏等等。我建议你在做生产环境时，至少加上插件健康检查、异常重启、性能监控这几个机制。别问我怎么知道的——都是血泪教训。