100 TensorRT Puglin开发

1. 核心原理

• 定义计算逻辑：告诉 GPU 如何执行这个特殊层的计算。
• 集成到 TensorRT：让 TensorRT 在推理时调用你的计算逻辑。

2. 开发步骤

1. 创建一个插件类，继承TensorRT的 IPluginV2 接口，实现关键函数

#include <NvInfer.h>

class MyPlugin : public nvinfer1::IPluginV2 {
public:
    // 必须实现的函数
    const char* getPluginType() const override { return "MyPlugin"; }
    const char* getPluginVersion() const override { return "1"; }
    int getNbOutputs() const override { return 1; }  // 输出张量数量
    nvinfer1::Dims getOutputDimensions(int index, const nvinfer1::Dims* inputs, int numInputDims) override {
        // 定义输出张量的形状（例如和输入形状相同）
        return inputs[0];
    }
    // 其他函数（序列化、反序列化、初始化等）...
};

2. 实现计算逻辑（核心）

在自定义插件类，enqueue函数编写CUDA核函数，定义此算子的数据处理步骤

int MyPlugin::enqueue(
    int batchSize, 
    const void* const* inputs,  // 输入数据指针数组
    void* const* outputs,        // 输出数据指针数组
    void* workspace,             // 临时显存空间
    cudaStream_t stream          // CUDA 流
) {
    // 调用 CUDA 核函数（例如逐元素加法）
    const float* input = static_cast<const float*>(inputs[0]);
    float* output = static_cast<float*>(outputs[0]);
    int numElements = batchSize * inputDims.d[0] * inputDims.d[1] * ...; // 计算总元素数
    myKernel<<<gridSize, blockSize, 0, stream>>>(input, output, numElements);
    return 0;
}

// 示例 CUDA 核函数（每个线程处理一个元素）
__global__ void myKernel(const float* input, float* output, int numElements) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < numElements) {
        output[idx] = input[idx] + 1.0f;  // 自定义计算（例如每个元素加1）
    }
}

3. 注册插件

告诉TensorRT自己实现了某个算子

class MyPluginCreator : public nvinfer1::IPluginCreator {
public:
    const char* getPluginName() const override { return "MyPlugin"; }
    const char* getPluginVersion() const override { return "1"; }
    nvinfer1::IPluginV2* createPlugin(const char* name, const nvinfer1::PluginFieldCollection* fc) override {
        return new MyPlugin();
    }
    // 其他函数（反序列化等）...
};

// 注册插件
REGISTER_TENSORRT_PLUGIN(MyPluginCreator);

4. 编译成动态库

# 示例编译命令（Linux）
g++ -shared -o libmyplugin.so my_plugin.cpp -I/path/to/tensorrt/include -L/path/to/tensorrt/lib -lnvinfer

4. ONNX转TensorRT，TensorRT自动调用注册插件

import tensorrt as trt

# 加载插件库
trt.init_libnvinfer_plugins(trt.Logger(), "")
with open("model.onnx", "rb") as f:
    engine = runtime.deserialize_cuda_engine(f.read())

---

补充：
1. TensorRT原理补充
TensorRT流程：

• 将模型（例如ONNX）转化为TensorRT内部计算图（由层节点和数据流边组成）
• 优化计算图
• 生成引擎：将优化后计算图编译成二进制文件，供推理时直接调用。

关于IPluginV2类

• TensorRT 的计算图由多个“层”（Layer）组成，每个层必须遵守统一的接口规范。
• IPluginV2 就是这个规范，它定义了所有层必须实现的功能，例如：
  • 计算逻辑：enqueue() 函数中实现 GPU 核函数
  • 内存管理：configurePlugin() 定义输入输出内存需求。
  • 序列化与反序列化：serialize() 和 deserialize() 用于保存和加载模型。
  • 形状推导：getOutputDimensions() 确定输出张量形状。

关于Creator 类的作用

• IPluginCreator 是插件的“工厂”，负责两件事：
  • 生产插件实例：当 TensorRT 解析到自定义层时，调用 createPlugin() 创建插件对象。
  • 注册插件信息：告诉 TensorRT 插件的名称、版本、参数格式，方便后续查找和调用。
• TensorRT 在解析模型时，可能需要在不同阶段（如构建引擎、反序列化引擎）动态创建插件。Creator 类通过统一的接口管理插件的生命周期，实现“即插即用”。

2. TensorRT版本是否和GPU型号相关？在一个型号转换的TensorRT文件，在另一个GPU型号可以直接使用吗？

• TensorRT 生成的推理引擎文件（.engine）与 GPU 型号直接相关，TensorRT 在优化模型时，会根据 GPU 的计算能力生成特定的内核代码。在 A 型号 GPU 上生成的 TensorRT 引擎文件，若直接拿到 B 型号 GPU 上使用，可能无法运行，除非两者满足特定条件。