Sparse4Dv3 部署到 TensorRT-（1）

在开始正式工作前，先学习几个软件库的使用：

• cuda代码库
• pytorch 的自定义算子
• 设计支持onnx的pytorch算子
• 导出onnx
• 设计支持tensorrt的算子库
• 用trt将onnx导出为tensorrt引擎文件
• 用c++ api 使用onnx和自定义算子

首先从了解从CUDA算子开发到ONNX导出，再到TensorRT插件集成的完整流程。

CUDA算子开发与部署全流程指南

1、 CUDA算子实现

首先需要编写CUDA核函数（Kernel）作为算子的核心计算单元，然后使用C++进行封装以便Python调用。

// src/my_op/my_op_cuda.cu
#include <cstdio>#define THREADS_PER_BLOCK 256
#define DIVUP(m, n) ((m + n - 1) / n)// CUDA核函数：实现元素级加法
__global__ void my_op_kernel(const float* input1, const float* input2, float* output, int n) {int idx = blockIdx.x * blockDim.x + threadIdx.x;if (idx < n) {output[idx] = input1[idx] + input2[idx]; // 这里可以是更复杂的计算}
}// 启动函数：配置线程网格和块
void my_op_launcher(const float* input1, const float* input2, float* output, int n) {dim3 blocks(DIVUP(n, THREADS_PER_BLOCK));dim3 threads(THREADS_PER_BLOCK);my_op_kernel<<<blocks, threads>>>(input1, input2, output, n);
}

// src/my_op/my_op.cpp
#include <torch/extension.h>
#include <torch/serialize/tensor.h>// 输入检查宏
#define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x, " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x, " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)// 声明CUDA启动函数
void my_op_launcher(const float* input1, const float* input2, float* output, int n);// Python接口函数
void my_op_gpu(torch::Tensor input1, torch::Tensor input2, torch::Tensor output) {CHECK_INPUT(input1);CHECK_INPUT(input2);CHECK_INPUT(output);const float* input1_ptr = input1.data_ptr<float>();const float* input2_ptr = input2.data_ptr<float>();float* output_ptr = output.data_ptr<float>();int n = input1.numel();my_op_launcher(input1_ptr, input2_ptr, output_ptr, n);
}// 使用pybind11封装为Python模块
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {m.def("forward", &my_op_gpu, "My custom operator (CUDA)");
}

2、使用setup.py编译安装

创建setup.py文件来编译CUDA算子并安装为Python包：

# setup.py
from setuptools import setup, find_packages
from torch.utils.cpp_extension import BuildExtension, CUDAExtensionsetup(name='my_custom_ops',  # 包名称version='0.1.0',       # 版本号author='Your Name',    # 作者packages=find_packages(),ext_modules=[CUDAExtension(name='my_op',  # 模块名sources=['./src/my_op/my_op.cpp','./src/my_op/my_op_cuda.cu',],extra_compile_args={'cxx': ['-O3'], 'nvcc': ['-O3', '--use_fast_math']}),],cmdclass={'build_ext': BuildExtension},install_requires=['torch',],
)

使用以下命令安装：

pip install -e .

3、PyTorch集成与自动求导

创建PyTorch自动求导函数以便集成到神经网络中：

# ops/ops_py/my_op.py
import torch
from torch.autograd import Function
import my_op  # 导入编译的CUDA算子class MyOpFunction(Function):@staticmethoddef forward(ctx, input1, input2):# 确保输入在GPU上且内存连续input1 = input1.contiguous()input2 = input2.contiguous()# 准备输出张量output = torch.zeros_like(input1)# 调用CUDA算子my_op.forward(input1, input2, output)# 保存用于反向传播的中间结果ctx.save_for_backward(input1, input2)return output@staticmethoddef backward(ctx, grad_output):# 获取前向传播保存的张量input1, input2 = ctx.saved_tensors# 计算梯度（这里需要根据实际运算实现）grad_input1 = grad_output.clone()grad_input2 = grad_output.clone()return grad_input1, grad_input2# 创建易用的函数接口
my_op_func = MyOpFunction.apply

4、适配 ONNX 导出自定义算子

为了将自定义算子导出到ONNX，需要实现符号函数[symbolic]：

# ops/ops_py/my_op.py
class MyOpFunction(Function):# ... 前面的forward和backward方法 ...@staticmethoddef symbolic(g, input1, input2):# 定义ONNX导出时的符号# 这里将自定义算子映射为ONNX的Add操作作为示例# 对于真正自定义的算子，可能需要使用自定义算子名称return g.op("Add", input1, input2)# 如果是真正的自定义算子，可能需要这样定义：# return g.op("MyCustomOp", input1, input2, #             attribute1_f=1.0,  # 浮点属性#             attribute2_s="value")  # 字符串属性# ONNX导出示例
def export_onnx():# 创建示例输入dummy_input1 = torch.randn(10, 10).cuda()dummy_input2 = torch.randn(10, 10).cuda()# 导出模型torch.onnx.export(MyOpFunction.apply,(dummy_input1, dummy_input2),"my_custom_op.onnx",opset_version=11,input_names=['input1', 'input2'],output_names=['output'],dynamic_axes={'input1': {0: 'batch_size'},'input2': {0: 'batch_size'},'output': {0: 'batch_size'}})

5、TensorRT插件开发

对于TensorRT插件，需要实现IPluginV2DynamicExt接口：

// tensorrt_plugin/MyCustomPlugin.h
#include <NvInfer.h>
#include <NvInferPlugin.h>
#include <cuda_runtime_api.h>
#include <string>
#include <vector>class MyCustomPlugin : public nvinfer1::IPluginV2DynamicExt {
public:MyCustomPlugin(const std::string& name);MyCustomPlugin(const void* data, size_t length);~MyCustomPlugin() override;// IPluginV2DynamicExt 方法nvinfer1::IPluginV2DynamicExt* clone() const noexcept override;nvinfer1::DimsExprs getOutputDimensions(int outputIndex, const nvinfer1::DimsExprs* inputs, int nbInputs,nvinfer1::IExprBuilder& exprBuilder) noexcept override;bool supportsFormatCombination(int pos, const nvinfer1::PluginTensorDesc* inOut,int nbInputs, int nbOutputs) noexcept override;void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in, int nbInputs,const nvinfer1::DynamicPluginTensorDesc* out, int nbOutputs) noexcept override;size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs, int nbInputs,const nvinfer1::PluginTensorDesc* outputs, int nbOutputs) const noexcept override;int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,const nvinfer1::PluginTensorDesc* outputDesc,const void* const* inputs, void* const* outputs,void* workspace, cudaStream_t stream) noexcept override;// IPluginV2Ext 方法nvinfer1::DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes,int nbInputs) const noexcept override;// IPluginV2 方法const char* getPluginType() const noexcept override;const char* getPluginVersion() const noexcept override;int getNbOutputs() const noexcept override;int initialize() noexcept override;void terminate() noexcept override;size_t getSerializationSize() const noexcept override;void serialize(void* buffer) const noexcept override;void destroy() noexcept override;void setPluginNamespace(const char* pluginNamespace) noexcept override;const char* getPluginNamespace() const noexcept override;private:std::string mPluginNamespace;std::string mPluginName;
};// 插件创建器
class MyCustomPluginCreator : public nvinfer1::IPluginCreator {
public:MyCustomPluginCreator();~MyCustomPluginCreator() override = default;const char* getPluginName() const noexcept override;const char* getPluginVersion() const noexcept override;const nvinfer1::PluginFieldCollection* getFieldNames() noexcept override;nvinfer1::IPluginV2* createPlugin(const char* name,const nvinfer1::PluginFieldCollection* fc) noexcept override;nvinfer1::IPluginV2* deserializePlugin(const char* name,const void* serialData, size_t serialLength) noexcept override;void setPluginNamespace(const char* pluginNamespace) noexcept override;const char* getPluginNamespace() const noexcept override;private:static nvinfer1::PluginFieldCollection mFC;static std::vector<nvinfer1::PluginField> mPluginAttributes;std::string mPluginNamespace;
};

// tensorrt_plugin/MyCustomPlugin.cu
#include "MyCustomPlugin.h"
#include <cstring>// 实现CUDA核函数
__global__ void my_custom_op_kernel(const float* input1, const float* input2, float* output, int n) {int idx = blockIdx.x * blockDim.x + threadIdx.x;if (idx < n) {output[idx] = input1[idx] + input2[idx]; // 实际计算}
}// 实现插件方法
int MyCustomPlugin::enqueue(const nvinfer1::PluginTensorDesc* inputDesc,const nvinfer1::PluginTensorDesc* outputDesc,const void* const* inputs, void* const* outputs,void* workspace, cudaStream_t stream) noexcept {int n = inputDesc[0].dims.d[0]; // 获取批量大小// 调用CUDA核函数dim3 blocks((n + 255) / 256);dim3 threads(256);my_custom_op_kernel<<<blocks, threads, 0, stream>>>(static_cast<const float*>(inputs[0]),static_cast<const float*>(inputs[1]),static_cast<float*>(outputs[0]),n);return cudaGetLastError();
}// 其他必要方法的实现...
// 注意：这里需要实现头文件中声明的所有方法

6、完整工作流集成

将上述所有组件集成到一个完整的工作流中：

# test_integration.py
import torch
from ops.ops_py.my_op import MyOpFunction, export_onnx
import tensorrt as trt# 测试CUDA算子
def test_cuda_operator():input1 = torch.randn(1024, device='cuda')input2 = torch.randn(1024, device='cuda')output = MyOpFunction.apply(input1, input2)print("CUDA算子测试成功!")print(f"输入1: {input1.cpu().numpy()[:5]}")print(f"输入2: {input2.cpu().numpy()[:5]}")print(f"输出: {output.cpu().numpy()[:5]}")# 导出ONNX模型
def test_onnx_export():export_onnx()print("ONNX导出成功!")# 验证TensorRT集成
def test_tensorrt_integration():# 这里需要先编译TensorRT插件并注册# 然后加载ONNX模型并构建TensorRT引擎print("TensorRT插件需要单独编译和注册")if __name__ == "__main__":test_cuda_operator()test_onnx_export()test_tensorrt_integration()

步骤	关键挑战	解决方案
CUDA算子	内存管理和线程优化	使用CUDA最佳实践和性能分析工具
PyTorch集成	自动求导实现	正确实现forward和backward方法
ONNX导出	符号映射	实现symbolic方法定义ONNX操作
TensorRT插件	接口复杂性	完整实现IPluginV2DynamicExt所有方法

7、TensorRT插件的两种范式

你提到的两种方法是 TensorRT 中使用自定义插件的常见方式。下面我为你分析它们的原理、差异和适用场景。

两种方法的工作原理

方法	命令/代码示例	核心原理
1. trtexec 命令行	trtexec --onnx=model.onnx --plugins=libmyplugin.so	通过 --plugins 参数在构建阶段预加载包含自定义插件的共享库文件（.so）。TensorRT 的 ONNX 解析器在解析网络时，会利用已注册的插件来处理相应的自定义算子。
2. C++ API 编程	在代码中显式调用 getPluginRegistry()->registerCreator() 并创建 IBuilder	在你的 C++ 应用程序中，在解析 ONNX 模型之前，手动将自定义插件的 IPluginCreator 实例注册到 TensorRT 的全局插件注册表中。之后，ONNX 解析器遇到同名算子时，会自动发现并使用已注册的插件创建者来创建插件实例。

性能差异分析

通常情况下，这两种方法生成的最终 TensorRT 引擎在推理性能上不会有差异。

性能差异主要存在于构建阶段（即引擎生成过程），而不是推理阶段。原因在于：

1. 核心计算相同：无论通过哪种方式注册，最终引擎中运行的插件代码都是同一个 CUDA 核函数。插件的 enqueue() 方法实现决定了推理时的计算效率，而注册方式不影响这部分。
2. 构建过程可能略有不同：
   • trtexec：是一个封装好的工具，其内部流程固定，可能会包含一些额外的初始化或检查步骤。
   • C++ API：由于是你自己的程序，构建流程更精简、直接，理论上构建过程可能稍快一些。但这个时间差异与模型转换的总时间相比通常微不足道。

因此，选择哪种方法主要基于便利性和集成需求，而非性能考量。

无论哪种方法，都需要确保：

1. 插件注册时机：插件必须在 TensorRT 解析 ONNX 模型之前完成注册，否则解析器将无法识别自定义算子并报错。
2. 算子名称匹配：ONNX 模型中的算子名称（node.op_type()）、版本和命名空间必须与你编写的插件创建者所提供的完全一致，否则 TensorRT 无法将两者匹配起来。
3. 插件实现：对于通过 FallbackPluginImporter 注册的插件，其构造函数通常只能获取在 ONNX 模型中定义为属性（Attribute）的参数。而权重（Weight）信息会作为图层的输入（Input）传递，一些网络结构信息（如输入尺寸、卷积核大小等）可能需要在 configurePlugin 或 enqueue 阶段动态获取。
4. 环境依赖：确保运行环境中包含插件库的所有依赖项（如特定版本的 CUDA、cuDNN 或其他动态库）。