ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

TensorRT与Plugin开发实战:从PyTorch到高效推理引擎的完整指南

TensorRT与Plugin开发实战:从PyTorch到高效推理引擎的完整指南 干了这么多年模型部署我越来越觉得TensorRT、Plugin这类词已经被很多人挂在嘴边但真正能把一条模型从PyTorch权重稳稳当当变成线上可用的TensorRT引擎、甚至被逼到要自己手写Plugin的人其实少之又少。这篇文章是我在“模型部署与推理优化”系列里的第六篇专门聊聊TensorRT和Plugin开发这件事。如果你是刚接触模型部署想搞明白pt文件怎么转换TensorRT、FP16和INT8到底怎么选如果你已经在用trtexec跑过几个模型但遇到自定义算子不支持、精度对不上、性能拉不满这些破事或者你想看看FastSAM这类模型在C环境里怎么落地——那这篇文章基本就是写给你看的。内容会从原理讲到实操再讲到我实际踩过的坑尽量把那些文档里不会明说的细节都翻出来。1. 为什么总说TensorRT是“部署必选项”而不是“可选项”很多刚入门的朋友会问PyTorch模型直接跑不也挺好吗为什么非要折腾TensorRT这个问题我每次都会被问到答案其实就藏在推理优化的核心矛盾里——训练框架的重心是灵活性部署框架的重心是极致效率两者天生就不一样。1.1 TensorRT到底优化了什么从一轮卷积说起我习惯用一个类比来解释TensorRT做的事普通推理框架像一个小作坊每来一个订单一个层都要单独找师傅、单独开一次工TensorRT更像一条流水线它先把整个订单拆开看一遍把能合并的工序合并、能提前准备的物料提前备好然后流水线一转效率自然就上来了。具体到技术层面TensorRT主要做了这几件事。第一是层融合Layer Fusion。以最常见的Conv BN ReLU为例如果按原始框架的逻辑跑每一层都要读写一遍显存三层的显存开销和kernel启动开销是实打实的。TensorRT会把Conv和BN的权重提前融合成新卷积权重再把ReLU合并到卷积的激活函数里最后变成了一个Conv算子。显存访问少了两轮kernel也少启动两次性能提升是立竿见影的。类似还有Concat Elementwise、QKV合并这类操作都是它做融合的重灾区。第二是内核自动调优Kernel Auto-Tuning。TensorRT为同一个算子准备了大量CUDA kernel实现比如卷积就有针对不同tile大小、不同vectorized load宽度、不同算法Implicit GEMM、Winograd、FFT等的版本。构建引擎的时候它会用你设定的profile反复benchmark选出一个当前GPU上最快的组合。这也是为什么同一个ONNX模型在不同显卡上构建出的engine性能表现可能差很多——因为它针对你的卡做了专门的调优。第三是精度校准与低精度推理。FP16很多人觉得只是“把float换成半精度”但实际工程里远没那么简单。INT8更是需要你用校准数据去统计每一层的激活值分布算出合适的scale因子。这块做得好不好直接决定模型精度损失能不能控制在可接受范围内。我见过不少人一开INT8精度就崩多数情况不是TensorRT不行而是校准数据集没选对、或者某些层对量化过于敏感。第四是显存复用和内存池管理。TensorRT会在构建期分析整个计算图把生命周期不重叠的中间张量分配到同一块显存上减少显存碎片和峰值占用。这对边缘设备和小显存显卡来说非常关键。这些优化叠在一起实际效果有多明显以我最近在Jetson Orin上部署的一个检测模型为例原始PyTorch用FP32推理单帧大概要38ms转成TensorRT FP16后掉到11ms再上INT8加校准能压到7ms出头。这个差距不是靠换硬件能追回来的这也是为什么TensorRT成了NVIDIA平台上部署绕不开的一环。1.2 什么场景该用TensorRT什么场景反而不要用虽然TensorRT很强但我还是要泼一盆冷水它只服务于NVIDIA GPU平台。你的部署目标如果是树莓派5、RK3588这类ARM设备或者纯CPU服务器TensorRT再强也跟你没关系。这些场景该看的是ONNX Runtime、NCNN、OpenCV DNN或者RKNN它们各自有各自的优化思路把TensorRT硬搬过去只会给自己添堵。再一个如果你的模型还在频繁迭代今天改个结构、明天换个算子那也没必要急着转TensorRT。Engine构建本身就要花时间而且网络结构一变就得重新构建这种场景先用ONNX Runtime顶着反而更灵活。等模型结构稳定了、要上线压测了再切到TensorRT做最终优化这是我认为比较理性的节奏。还有个容易踩的误区TensorRT不是所有算子都支持。虽然官方算子覆盖面已经很大但Transformer里某些奇葩的mask组合、一些自定义算子、或者新paper里刚出的结构经常会在转engine时报“unsupported”错误。这时候你有三条路改网络结构绕开、用plugin扩展、或者退回ONNX Runtime。关于Plugin怎么开发后面有一章专门讲这里先记住一点——遇到不支持不等于死路TensorRT留了Plugin这个口子。2. 上手准备从PyTorch权重到可跑的TRT引擎我在帮别人排查部署问题的时候发现很多人连TensorRT安装都没搞定就开始问Plugin开发这有点本末倒置。先把基础链路走通比什么都有用。这一章我讲的就是我自己最常用、也最推荐大家走的链路PyTorch权重 → ONNX → TensorRT Engine。2.1 环境准备别让版本问题毁掉你的一天TensorRT最折磨人的不是它本身多难而是版本匹配问题。CUDA、cuDNN、TensorRT、显卡驱动这四者之间必须互相兼容有一个对不上编译直接报一堆莫名其妙的错。我踩过的坑就不止一次驱动版本太低cuda runtime初始化失败TensorRT 8.5的engine在8.6里加载报版本不匹配ONNX的opset版本太高TRT的parser解析不了。所以我的建议很简单如果你不是有特殊定制需求直接拉NVIDIA官方NGC的TensorRT容器镜像来用。NGC容器已经帮你把CUDA、cuDNN、TensorRT的匹配版本全部配好了你只需要保证宿主机的显卡驱动足够新就行。比如你换了新显卡想用最新的CUDA特性那大概率需要升级驱动这个查一下NVIDIA驱动支持矩阵就行了。如果确实要在宿主机上安装我建议用tar包方式代替deb包。tar包解压就能用路径可以自己控制不会污染系统环境。官网下载的时候注意选择和你CUDA版本匹配的TensorRT版本比如CUDA 12.x对应TensorRT 8.6和10.x系列。装完之后可以用一个极小测试验证一下随便导出一个ONNX模型trtexec先转一下能转成功说明基本环境没问题。# 以TensorRT 10.x为例安装tar包后设置环境变量 export TRT_RELEASE/opt/TensorRT-10.0.0.6 export LD_LIBRARY_PATH$TRT_RELEASE/lib:$LD_LIBRARY_PATH export PATH$TRT_RELEASE/bin:$PATH # 验证安装用自带的engine样例 trtexec --help | head -20提示如果trtexec缺libnvinfer.so等动态库用ldd $(which trtexec)看一下缺哪个把对应路径加到LD_LIBRARY_PATH里就行。这里看不到缺库就正常跑。2.2 导出ONNX时最容易踩的三个坑ONNX是中间格式导出这一步的质量直接影响后面转engine的成功率。我总结了三个高频坑都是我自己和身边同事反复踩过的。第一个坑是动态轴设置不完整。很多模型输入是NCHW格式部署时才需要动态batch或者动态分辨率。你用torch.onnx.export导出时必须明确告诉它哪些维度是动态的。常见写法是这样import torch dummy_input torch.randn(1, 3, 640, 640) dynamic_axes { input: {0: batch, 2: height, 3: width}, } torch.onnx.export( model, dummy_input, model.onnx, opset_version17, input_names[input], output_names[output], dynamic_axesdynamic_axes, )如果你漏了dynamic_axes后面在TensorRT里想支持动态尺寸就得重新导出白折腾一遍。第二个坑是opset版本和算子兼容性。ONNX的opset版本一直在更新TensorRT对不同opset版本的支持程度不一样。我一般建议导出时opset用17或者18兼容性和新特性之间比较均衡。遇到某个算子解析不了先看看是不是opset太高或者太低导致的。另外有些PyTorch算子比如某些index_put、einsum导出的ONNX结构非常绕TensorRT解析起来费劲甚至直接失败这时候可能要考虑在导出前对模型做一点等价的算子替换。第三个坑是自定义算子。如果你的模型里有PyTorch原生算子无法表达的op导出ONNX时会报错或者生成一个不完整的图。常见做法是先把自定义op在PyTorch里注册成ONNX的自定义符号register_onnx_op或者在导出时用torch.onnx.is_onnx_supported提前检查哪一层不支持再针对性地处理。这块如果处理不好后面到了TensorRT大概率也会卡在同一个算子上——那就要跳到第三章Plugin开发了。2.3 trtexec转engine参数选型与实际操作导出ONNX成功之后转engine反而是最简单的一步。我基本都用trtexec这个工具它足够强大还能顺带做性能benchmark。一个基础的转换命令长这样trtexec \ --onnxmodel.onnx \ --saveEnginemodel.engine \ --fp16 \ --workspace4096这里几个参数要说清楚。--fp16表示启用FP16精度。如果模型对精度不太敏感这一步能让性能直接上一个台阶。但注意如果你的网络里有容易溢出的操作比如大的激活值、LayerNorm、Softmax之前的值FP16可能会让精度明显下降后面5.1节我会详细说怎么排查。--workspace是给TensorRT构建engine时用的临时显存上限单位是MB。它影响的是构建期的优化空间不是运行期。历史版本里这个参数出现过改名新版本用--memPoolSize。如果你发现同一个模型构建出来的engine性能不太理想可以试着把这个值调大一点让TensorRT有更多空间去尝试更好的kernel方案。如果要做INT8还需要提供校准数据。trtexec里可以指定一个校准图片目录trtexec \ --onnxmodel.onnx \ --saveEnginemodel_int8.engine \ --int8 \ --calib/path/to/calibration/images对于动态shape的模型转engine时必须同时指定三个profile最小shape、最优shape、最大shape。这三个值直接决定运行时能接受的输入范围也影响kernel调优策略。比如目标输入是1到8的batch分辨率变化范围在320到640之间trtexec \ --onnxmodel.onnx \ --saveEnginemodel_dynamic.engine \ --fp16 \ --minShapesinput:1x3x320x320 \ --optShapesinput:4x3x640x640 \ --maxShapesinput:8x3x640x640optShapes是最关键的一个它告诉TensorRT你预期的常见shape是什么TensorRT会优先为这个shape选kernel。如果你平时推理时输入基本是固定尺寸比optShapes设置得偏大或偏小都会导致性能打折扣。这里的“性能”不是指报错而是kernel选得不精准跑得比该有的速度慢——这是很多人用了动态shape之后抱怨“变慢了”的常见原因。转完之后用--loadEngine配合--shapes就能做性能测试方便快速确认engine在不同shape下的延迟表现。3. Plugin开发的完整拆解从Why到How如果说转engine是体力活那Plugin开发就是真正的技术活了。会写Plugin的人在团队里通常就是那个“搞不定找他”的角色。我写这一章不想只堆API更想把思路讲透什么时候该写、怎么写、写的时候哪些坑是躲不开的。3.1 什么时候才需要写Plugin我在1.2节里提过TensorRT不是所有算子都支持。但“算子在ONNX里存在”和“值得为它写Plugin”是两码事。我总结出三个真正需要写Plugin的典型情况。第一种TensorRT不支持某个算子。比如一些论文里刚出来的特殊注意力结构、可变长mask组合、或者自定义的量化反量化逻辑。这种情况下不写Plugin就卡死了属于被迫营业。第二种算法层面有明显可融合的空间。比如一个复杂的block里多个算子可以合并成一个自定义kernel融合之后能大幅减少中间张量的显存读写。这种Plugin写出来价值最高因为它不只是“补漏”而是主动优化。我在FastSAM部署时就用过一个类似思路——把一些interpolation和mask计算融合进一个自定义kernel减少显存往返延迟直接降了30%。第三种推理链路里有一些TensorRT抽象不了的特殊逻辑。比如需要读取外部数据、需要和某套硬件库做集成、或者推理时需要同步执行一些复杂的后处理步骤。这些逻辑写在Plugin里能跟engine的stream对齐比在host端单独做要高效得多。反过来如果只是个别算子不支持但你的模型结构允许绕开比如把自定义算子拆成几个TensorRT支持的基础算子组合那就没必要写Plugin——毕竟Plugin要维护、要跨平台编译、版本升级还可能炸成本真的不低。3.2 一个Plugin的骨架长什么样TensorRT的Plugin接口版本演进过好多次老代码里常见的是IPluginV2、IPluginV2IOExt、IPluginV2DynamicExt。我的建议是除非你明确知道自己只处理静态shape否则直接上IPluginV2DynamicExt。它支持动态输入形状是现在的主流做法也更容易适配多样化的部署场景。一个典型的Plugin类大概长这样class MyCustomPlugin : public nvinfer1::IPluginV2DynamicExt { public: // 返回插件名称、版本、命名空间 const char* getPluginType() const override; const char* getPluginVersion() const override; // 动态shape下计算输出张量的维度 DimsExprs getOutputDimensions( int outputIndex, const DimsExprs* inputs, int nbInputs, IExprBuilder exprBuilder) override; // 判断给定输入格式组合是否支持 bool supportsFormatCombination( int pos, const PluginTensorDesc* inOut, int nbInputs) override; // 配置插件确认输入输出格式等 void configurePlugin( const DynamicPluginTensorDesc* in, int nbInputs, const DynamicPluginTensorDesc* out, int nbOutputs) override; // 初始化与销毁资源 int initialize() override; void terminate() override; // 核心推理逻辑在GPU stream上执行 int enqueue(const PluginTensorDesc* inputDesc, const PluginTensorDesc* outputDesc, const void* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) override; // 序列化相关网络模型保存为engine时需要 size_t getSerializationSize() const override; void serialize(void* buffer) const override; // 克隆插件用于多context场景 IPluginV2DynamicExt* clone() const override; };enqueue是整个插件的灵魂里面要在指定的CUDA stream上启动kernel而且必须是非阻塞的——直接调用kernel函数就行千万别在enqueue里做同步等待不然整个推理流程的异步流水就废了。int MyCustomPlugin::enqueue(...) { // 从输入输出desc里取到当前batch、shape等信息 int batch inputDesc-dims.d[0]; int channels inputDesc-dims.d[1]; // 直接调用自定义CUDA kernel走传进来的stream myCustomKernelgrid, block, 0, stream( inputs[0], outputs[0], batch, channels, ...); return 0; }这里有个细节enqueue返回0表示成功。有些新手会在这里写cudaStreamSynchronize(stream)想确保kernel执行完结果把性能彻底毁了还会引入死锁风险。记住kernel启动是异步的同步交给TensorRT的context去管.3.3 序列化与动态Shape两个绕不开的难点写Plugin时有两个设计难点几乎每个人都会遇到。第一个是序列化和反序列化。TensorRT构建engine时会把网络结构、权重以及每个plugin需要的参数都序列化成二进制数据。你要在serialize里把自己的参数写进buffer在PluginCreator的deserializePlugin里再读出来。这个buffer是按字节排布的所以顺序和sizeof必须完全一致建议在buffer开头写一个版本号方便以后扩展字段——不然你改了插件结构老的engine就全部作废且你根本不知道错在哪里。第二个难点是动态shape下的维度推导。getOutputDimensions要在没有具体数值的情况下用表达式算出输出维度例如输出宽等于输入宽减掉padding。写法是用exprBuilder的operation(DimensionOperation::kSUB, ...)这类方法来表达维度运算。这里很容易写错我建议把每个分支的维度推导都测一遍千万不要只测一个shape就以为没问题。再补充一个坑如果Plugin中有权重参数记得在initialize里拷贝到GPU显存在terminate里释放。有些人在构造函数里就开辟显存、在析构函数里释放看起来没错但engine在反序列化时需要先创建Plugin再拷贝权重这中间的状态流转很容易出问题最后表现就是同一份代码在一个场景正常、另一个场景崩溃。3.4 注册、编译与调试让TensorRT认你的插件写好了类还不够TensorRT得能发现你的插件。这需要两步继承IPluginCreator实现一个creator类然后用宏注册。class MyCustomPluginCreator : public nvinfer1::IPluginCreator { public: const char* getPluginName() const override { return MyCustom; } const char* getPluginVersion() const override { return 1; } nvinfer1::IPluginV2* createPlugin( const char* name, const nvinfer1::PluginFieldCollection* fc) override; nvinfer1::IPluginV2* deserializePlugin( const char* name, const void* serialData, size_t serialLength) override; }; REGISTER_TENSORRT_PLUGIN(MyCustomPluginCreator);注意宏的版本差异。老版本用REGISTER_PLUGIN新版8.x后期和10.x用REGISTER_TENSORRT_PLUGIN。这个宏会在plugin库被加载时自动注册所以你的插件库编译后要在推理程序里显式加载它。加载方式有两种一是用dlopen手动加载so文件再把so路径加到TensorRT搜索路径二是把so放到TensorRT的plugin目录下。我推荐第一种因为部署时对路径的可控性更强#include dlfcn.h void* handle dlopen(libmyplugins.so, RTLD_NOW); if (!handle) { std::cerr load plugin lib failed: dlerror() std::endl; return -1; } auto creator nvinfer1::getPluginRegistry()-getPluginCreator(MyCustom, 1); if (!creator) { std::cerr creator not found std::endl; return -1; }编译时链接TensorRT的头文件和库基本就能过但有几个细节我强调一下Plugin库最好用C11及以上标准编译新版TRT头文件对C版本有要求编译机器和运行机器的GCC版本别差太多否则会有ABI兼容问题能静态链接的依赖尽量静态链减少运行时动态库依赖。调试Plugin是很多人的噩梦因为直接跑trtexec或者推理程序程序崩了你根本不知道崩在哪。我的调试习惯是先在CPU端用一个简单的参考实现把逻辑跑通再在CUDA kernel里只放最简单的copy逻辑验证接口生命周期是对的然后再逐步把核心逻辑加回来。如果崩溃了先用compute-sanitizer旧版叫cuda-memcheck跑一遍它能精确定位到kernel里哪一行越界了比printf大法高效太多了。4. 实操现场C推理流程与典型模型部署复现理论讲了不少这章回到真正的部署实战。我挑了两个典型的场景来展开一个是C API调用engine做推理的完整流程很多人在Python里跑得很溜一换C就抓瞎另一个是结合FastSAM在C上部署时踩过的Plugin相关坑跟第三章的Plugin内容做个呼应。4.1 用Engine做推理C API的核心流程engine构建好之后部署侧的代码其实比较套路化。我以C为例走一遍完整流程。第一步反序列化engine。序列化好的engine文件是一堆二进制字节要先用initLibNvInferParsers初始化解析器库和创建runtime再把engine数据喂进去。代码骨架nvinfer1::IRuntime* runtime nvinfer1::createInferRuntime(logger); std::ifstream engineFile(model.engine, std::ios::binary); std::vectorchar engineData(std::istreambuf_iteratorchar(engineFile), {}); nvinfer1::ICudaEngine* engine runtime-deserializeCudaEngine(engineData.data(), engineData.size());第二步创建context并准备输入输出buffer。这里最常见的错误是buffer大小算错——TensorRT在动态shape下buffer必须按最大可能shape分配否则推理时容易越界。比如输入最大支持8x3x640x640那就按这个大小分配显存即使当前实际只推理1张图。nvinfer1::IExecutionContext* context engine-createExecutionContext(); // 设置当前推理的实际shape context-setInputShape(input, nvinfer1::Dims4{1, 3, 640, 640}); // 分配显存buffer统一按最大shape分配 void* inputBuffer; void* outputBuffer; cudaMalloc(inputBuffer, 8 * 3 * 640 * 640 * sizeof(float)); cudaMalloc(outputBuffer, maxOutputSize);第三步用CUDA stream异步跑推理。新版API里有个enqueueV3需要传入一组张量地址TensorAddress比老的enqueueV2传buffer数组更简洁。我习惯用stream做异步拷贝让H2D拷贝、kernel执行、D2H拷贝三者流式重叠起来这样吞吐能提升不少。cudaStream_t stream; cudaStreamCreate(stream); // 输入数据拷到GPU cudaMemcpyAsync(inputBuffer, hostInput.data(), actualBytes, cudaMemcpyHostToDevice, stream); // 推理核心调用 context-enqueueV3(stream); // 输出拷回CPU cudaMemcpyAsync(hostOutput.data(), outputBuffer, outputBytes, cudaMemcpyDeviceToHost, stream); // 最后同步 cudaStreamSynchronize(stream);有个小技巧如果推理线程和拷贝线程可以分开就用双stream一个做拷贝一个做计算吞吐能再上一个台阶。当然这属于优化范畴了先跑通再说。4.2 FastSAM的TensorRT实践要点FastSAM这类模型是典型的“结构看着不大部署起来头大”的模型。它包含ViT结构、mask解码、以及一些自定义的计算逻辑。在C环境下部署FastSAM的TensorRT版本我遇到了几个问题正好和前面的章节串起来。首先是ViT结构导出ONNX时attention里的reshape和transpose特别多TensorRT解析通常没问题但engine构建时间会很长而且kernel调优的搜索空间很大。我的做法是固定输入分辨率比如640x640减少动态shape带来的组合爆炸这样可以明显缩短构建时间运行时也更可控。然后是mask解码部分。FastSAM的post-processing包含很多自定义操作有些TensorRT支持得不好。我的方案是把能放上GPU的部分尽量保留在engine里把实在搞不定的部分比如某些拓扑相关的处理拆出来放host端做。虽然多了一次D2H拷贝但换来的是不用为几个小众算子写一堆Plugin权衡下来部署更稳、工期更短。第三是输出层的处理。ViT类模型输出的tensor往往很大如果你在ONNX导出时就把后续的降维和后处理算子加进去可以减少输出显存拷贝量。这一步对整体延迟影响不大但能降低显存带宽压力在边缘设备上效果很明显。4.3 树莓派5部署YOLOv5先分清GPU平台再谈TensorRT这个标题我想了很久怎么写因为很多人在树莓派5上部署自己训练的YOLOv5模型时会搜到TensorRT相关教程然后带着TensorRT的思路去折腾。这里一定要说清楚TensorRT是NVIDIA GPU的专属推理优化库树莓派5用的是Broadcom的SoC根本没有NVIDIA GPU所以TensorRT在树莓派5上压根跑不了。那树莓派5上部署YOLOv5该用什么常见选择是NCNN、ONNX Runtime和OpenCV DNN。我自己在树莓派5上试过YOLOv5s用ONNX Runtime配合ARM的CPU优化单帧推理大概在200ms到300ms这个量级具体取决于分辨率和线程数。如果你想追求性能NCNN在ARM CPU上通常有更好的卷积优化而且它还支持RGA等硬件加速接口如果设备有NPU的话值得优先尝试。如果你的部署目标是Jetson系列比如Jetson Orin Nano那才能用TensorRT而且能感受到CUDA核心带来的巨大加速。同样是YOLOv5s在Orin Nano上转成TensorRT FP16之后单帧可以跑到10ms以内这个差距就是硬件平台推理框架共同决定的。所以决策路径应该是先确认目标设备的计算单元是什么再选择推理框架。NVIDIA GPU → TensorRTARM CPU/GPU → NCNN或ONNX Runtime其他厂商NPU → 各自SDK比如RKNN。不要看着TensorRT跑得快就觉得所有设备都要用它框架选错了后面的努力基本白费。5. 常见问题与排查技巧实录部署这个行当花最多时间的从来不是“跑通”而是“排查”。这一章我把我自己踩过、以及帮别人排查过的高频问题整理成了一份实录每一类都给出排查路径而不是直接丢结论。这样你遇到问题的时候可以照着自己的实际情况去定位。5.1 精度对不上从API到算子的完整排查路径“FP16之后输出跟PyTorch对不上”是我收到的最常见的求助类型。这里要先界定清楚“对不上”是量级上的偏差还是完全错误的结果。如果是完全错误多半是网络结构解析出了问题或者Plugin的维度算错了如果有细微偏差那大概率是精度调优的问题。我一般按照这个顺序排查先用FP32构建一个engine不做任何低精度优化。如果FP32的输出和PyTorch对不上说明问题出在转换/解析环节先查ONNX导出和TensorRT解析别急着甩锅给精度。FP32能对上再开FP16。开了FP16后偏差变大先看模型里有没有对数值范围敏感的层比如LayerNorm、Softmax之前的大激活值、或者einsum累加。常见解法是在这些层前后插入clip或者强制某些层用FP32计算。如果开了INT8之后偏差继续变大先检查校准数据集。校准集必须能代表真实推理时的输入分布我见过有人用随机噪声做校准最后推理真实图片时精度崩得没法看。另外不同校准算法histogram、entropy、minmax对结果影响很大可以多试几种对比一下。排查工具方面TensorRT提供了Polygraphy这个神器。它能做逐层对比定位到底是哪一层开始精度产生偏差还可以自动跑多种精度和算法组合。有它在精度排查效率至少提升三倍强烈建议部署工具箱里常备。5.2 性能不达预期延迟比预期高先检查这四件事性能问题是另一大类。如果你的engine跑出来的延迟跟网上类似配置的差距很大先不要怀疑TensorRT不行按下面四个方向自查。第一动态shape的profile是否配置得当。很多人图省事只设了min和max没设optShapes。或者optShapes设得很大但实际推理时输入很小导致TensorRT选的kernel不是最优的。正确的做法是optShapes尽量贴近实际推理的常见shape。第二workspace/memPoolSize是否够大。构建engine时的临时显存上限会影响kernel选择范围设置太小的话TensorRT可能被迫选择保守但慢的kernel。这个值可以试着调大对比一下延迟变化。第三是否开启了CUDA Graph或者做了stream复用。TensorRT在Python下有个trt_cuda_graph的玩法C下也可以把推理用到的kernel序列包裹成CUDA Graph减少kernel启动开销。当你的单帧延迟已经比较低比如2ms以下时kernel启动耗时占比会很高CUDA Graph能有效改善。第四CPU和GPU之间的数据传输是否成了瓶颈。如果在边缘设备上PCIe带宽或者内存带宽有限H2D/D2H拷贝时间可能比推理本身还长这时要考虑把更多的前后处理也挪到GPU上或者用异步拷贝把传输时间和计算时间重叠起来。5.3 Plugin加载失败从注册到版本的全链路排查Plugin相关的问题我单独列一节因为它最容易让人陷入“代码没问题但就是跑不起来”的怪圈。常见的失败场景有几种一种是加载so时报符号找不到。这种情况基本是编译时的依赖没链全比如Plugin里用了OpenCV的库但推理程序里没有链接OpenCV。用ldd libmyplugins.so查看依赖把缺的库补上就行。另一种是getPluginCreator返回空指针。可能是so没有成功加载dlopen失败也可能是plugin的name和version跟你传入的不匹配。我调试时通常会在creator里加一行printf输出plugin name确认它到底有没有被注册进去。还有一种比较隐蔽老engine文件里的plugin序列化数据跟新版plugin代码不兼容。这种情况通常表现为反序列化时崩溃或结果错误我在3.3节里说过序列化buffer设计时一定要带版本号这样至少能及早报错而不是莫名其妙出错。最后强调一个ABI和运行时版本的问题。TensorRT本身是版本强相关的你用10.0编译的Plugin跑到TensorRT 8.6的环境里几乎必然出问题。线上部署环境尽量和开发环境保持一致包括TensorRT小版本、CUDA版本、以及GCC版本这是最简单也最容易被忽略的坑。5.4 大模型方向TensorRT-LLM与Docker部署的衔接虽然这篇文章的主角是通用TensorRT和Plugin但现在做部署的人很难不碰上大模型相关的话题。热搜词里出现了“onnx部署llm模型”、“docker部署vllm模型教程”——说明很多人正在用ONNX、Docker、vLLM这些东西做大模型推理。我之前写过一篇关于vLLM部署并正确配置资源的文章很多粉丝反馈说在官方docker镜像比如vllm/vllm-openai里部署后访问量一大就OOM或者响应卡顿。我的排查思路是完全按官方文档拉镜像、重启时会先看显卡的显存占用和CPU内存占用然后限制并发数再开--swap-space配置。中间的分析严格基于环境的实际资源和日志信息不依赖网络上的所谓“优化技巧”。在Linux本机Docker环境里用官方镜像部署时我用CRI和GPU拓扑相互印证把NVIDIA GPU的利用率、NUMA节点和CPU pinning配套调好实测在多并发场景下吞吐稳定了不少。当然部署完模型后的可视化也很重要这类知识在之前文章里详细写过这里就不再展开了。6. 我在实际部署中得到的几点体会这篇写完内容已经不少了最后分享一点我做模型部署这几年沉淀下来的体会不算是总结更像是个人的一些“手记”。第一个体会能不改模型就不改模型能少写Plugin就少写Plugin。很多人一看到TensorRT不支持某个算子第一反应就是写Plugin。但Plugin的维护成本是持续的——TensorRT版本升级、推理设备更换、甚至GCC版本变动都可能让你重新编译调试。我见过太多团队被一两个自研Plugin拖住每次升级都要陪着重新趟一遍坑。如果你能用已有算子组合绕开或者在导出ONNX之前对模型做一点等价调整往往能省下大量时间。只有在融合优化能带来显著收益、或者实在无法避开时写Plugin才是划算的选择。第二个体会精度和性能永远是一个工程权衡的问题不是非黑即白。FP16不行就上INT8试试INT8校准集不准就换算法。关键是你要建立一条可复现的对比链路把PyTorch、ONNX Runtime、TensorRT FP32/FP16/INT8的精度和延迟都测一遍心里有数。这样每次优化都有据可依而不是凭感觉调参。第三个体会部署的稳定性往往取决于你对手里工具链的熟悉程度而不是某个单一框架有多强。trtexec、Polygraphy、compute-sanitizer、NVIDIA Nsight这些工具每一样都值得花时间掌握。我在排查问题时80%的时间其实是在用工具定位问题真正动手改代码的时间很少。工具用得越熟解决问题的速度就越快。最后分享一个小技巧每次部署一个新模型我都会强制自己把构建命令、关键参数、踩过的坑记录下来哪怕只是给自己看的几行字。半年后你回来看这些记录会觉得当时的自己帮你省了巨大的时间。这个习惯也是我这整个“模型部署与推理优化”系列一直在坚持写作的原因。
返回列表