开头那两段先接着过去半年我一直在折腾一个老掉牙但绕不过去的问题TensorRT7 Onnx模型动态多batch怎么才能稳定跑起来。网上翻到的资料大多停留在“加个profile就好了”真到自己动手才发现光配置profile远远不够很多隐藏的坑会从模型图、插件、甚至解析器解析结果里冒出来。这篇文章就把我完整踩过的一轮坑拆开讲清楚适合正在用TensorRT7做ONNX模型部署、又被动态batch卡住的人看。1. 动态batch需求的来源与最初的失败现象1.1 为什么线上服务会盯上动态batch我的场景很简单一个ONNX导出的图像检测模型要封装成内部推理服务。线上请求量波动很大低谷期几十个并发高峰期能冲到几千。如果模型把batch固定在16低谷期就得凑够16张图才发一次推理延迟高得没法看如果固定在1高峰期GPU算力又严重浪费吞吐上不去。最理想的情况就是让batch跟着请求量走模型每次推理实际进来的图有几张就用几张——这就是动态batch。用ONNX Runtime验证这个需求毫无压力把输入维度设置成[batch, 3, height, width]batch直接传一个动态维度就能正常工作。初次跑通的时候我还暗自庆幸觉得这个需求应该挺简单结果一搬到TensorRT7上问题像连环雷一样炸开。1.2 第一次移植就翻车的现场记录先把最基础的环境交代一下TensorRT 7.2.2CUDA 11.1ONNX模型是一个包含预处理、主干网络、后处理输出的检测模型输入名是input输出名是detections。第一次尝试我照着常见教程的套路写了个转换脚本import tensorrt as trt TRT_LOGGER trt.Logger(trt.Logger.WARNING) builder trt.Builder(TRT_LOGGER) network builder.create_network(1 int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) parser trt.OnnxParser(network, TRT_LOGGER) with open(model.onnx, rb) as f: parser.parse(f.read()) config builder.create_builder_config() config.max_workspace_size 1 30 profile builder.create_optimization_profile() profile.set_shape(input, min(1, 3, 640, 640), opt(4, 3, 640, 640), max(16, 3, 640, 640)) config.add_optimization_profile(profile) engine builder.build_serialized_network(network, config)理论上这套代码就是标准答案显式batch模式、profile设置动态维度、opt选个中间值。但执行完后用engine.get_binding_shape(input)去查输入维度看到的居然是固定shape(1, 3, 640, 640)而不是我设置的四维动态范围。用context.set_binding_shape去改batch维度直接报错说这个维度不允许动态修改。更诡异的是我把opt从4改成8重新构建后输入shape又变成(8, 3, 640, 640)固定值。说明profile确实被读取了但最终生效的永远只有那个opt值动态维度根本没被保留下来。当时我一度怀疑是TensorRT7对ONNX解析器有某种限制后来又怀疑是我的模型结构太复杂让解析器在某些层上把动态信息丢掉了。2. 先搞清楚TensorRT7的batch机制EXPLICIT_BATCH和profile的坑2.1 implicit batch和explicit batch到底差在哪排查到这一步我被迫把TensorRT7的batch机制从头捋了一遍。TensorRT从很早的版本开始网络定义就默认走“隐式batch”模式。隐式batch的意思是网络里的张量维度是[C, H, W]第一个维度C代表通道batch不作为一个真正的维度出现在网络定义里而是在执行时通过context.set_batch_size单独指定。这种设计对CNN很友好因为大部分层对batch是逐元素独立操作的batch设成几本质上不影响每个样本的计算路径。但对ONNX模型来说问题就明显了ONNX的输入维度通常是[N, C, H, W]N是实打实的一个张量维度很多算子和shape变换都会用N做计算。解析器如果默认按隐式batch来解释就会把N这个维度抹掉或者干脆当成固定值处理。所以TensorRT7要求处理ONNX模型时显式开启EXPLICIT_BATCH也就是create_network时传入那个flag。开了之后batch才会作为网络张量的第一维参与解析。这个知识点不算冷门但很多人只是机械地加了flag并不知道它意味着批处理维在后面的所有API调用里都不再是“特殊维”它只是一个普通的动态张量维度。2.2 profile只是告诉引擎“可以跑多大”不等于解析器自动念对了batch位置接下来是第二个大坑profile的作用范围。create_optimization_profile里设置的min/opt/max逻辑上只是给优化器一个搜索范围让TensorRT在选kernel、定显存策略时知道动态维度的边界。它并不会自动修正网络内部对batch维的引用。我后来的实测发现如果网络里有层在解析时把N写死了比如某个Reshape的shape值直接是[16, 56, 56]那么解析器就会把上游传来的动态batch范围截断成固定值profile设置的动态区间到这一层就断了。更隐蔽的是有些层不会立刻报错它会在构建阶段悄悄把动态维度退化成opt值对应的固定维度最后产出的engine里输入shape就成了死数。我把两种模式下的行为整理成一个表方便你对照配置方式batch在engine里的表现能不能运行时动态改隐式batch set_batch_sizeengine不含batch维执行时单独指定可以改但只适合简单网络ONNX解析基本不能用显式batch 无profile所有动态维被固定成维度中的第一个数不能改显式batch profile理想情况是min~max范围只有在解析链路完整保留动态信息时才成立显式batch profile 网络里有层把N写死最终engine退化成opt固定值不能改这就是我前文遇到的情况2.3 一个关键函数getBindingDimensions和executeV2真正帮我确认问题出在哪一步的是这两个APIgetBindingDimensions和executeV2。在engine构建完成后逐个打印输入输出绑定的维度for i in range(engine.num_bindings): name engine.binding_name(i) shape engine.get_binding_shape(i) print(f{name}: {shape})在跑动态shape时执行不再是context.execute而是context.execute_async_v2或者context.execute_v2。动态输入维度要通过context.set_binding_shape告诉context这次推理实际用多大batch而且每个profile范围都必须提前设置。如果engine里的binding shape是死的set_binding_shape一调用就直接报错根本没有进入执行阶段。3. 复现与逐层定位问题根本不在profile设置3.1 排除法第一步最小网络验证既然怀疑链路太长我就先砍出一个最小网络来验证用一个只有ConvReLUConv的ONNX模型输入动态batch。用同一套转换脚本这次engine里的binding shape正常出现了动态范围(-1, 3, 640, 640)运行时把batch改成1、4、8、16都完全正常。这一下就证明了两件事TensorRT7本身对动态batch的解析机制没问题我的profile和EXPLICIT_BATCH用法也没问题那么问题只能出在真实模型的网络结构里。接下来就是用排除法逐层缩小范围。3.2 定位到网络结构里唯一的“不透明”部分自定义插件真实模型里有一堆自定义算子一个基于像素聚合的特殊算子为了跑TensorRT我只得用Plugin API封了一层。于是第二个实验很快设计出来把模型里的自定义算子和它前后连接的部分全部去掉用一堆标准卷积层重新连了一个简化网络再次转engine。结果动态范围又正常了。到这一步嫌疑基本100%落在插件上。打开TensorRT7里的自定义插件基类所有实现动态shape的插件都必须继承IPluginV2DynamicExt注意不是旧的IPluginV2后者根本不支持动态维度。我的插件当时继承的是旧版基类方法列表里压根没有getOutputDimensions的逐维度计算逻辑只是在初始化的时候把输出shape直接写死。解析器一看插件输出shape是固定的自然而然就把后续所有动态信息都掐断了。3.3 插件内部三个藏着固定batch的地方即使改用了IPluginV2DynamicExt插件里还有三个位置特别容易把batch写死我逐一检查并修正第一处getOutputDimensions。这个方法要根据输入shape推导输出shape正确的做法是从inputDims[0].d[0]取动态batch值而不是写死一个数字。我最初的代码里写的是outputDims[0].d[0] 16改完后变成Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override { Dims out; out.nbDims 4; out.d[0] inputs[0].d[0]; // 动态batch不动 out.d[1] inputs[0].d[1]; out.d[2] 640; // 其他维度可以写死 out.d[3] 640; return out; }第二处supportsFormatCombination。这个方法是告诉TensorRT插件支持哪些输入输出格式组合它本身不直接涉及shape但TensorRT7内部在构建优化时会拿它做形状推断。如果这个方法对某些非batch维度返回false解析器会被迫把shape固定。检查一下确认格式组合没问题。第三处enqueue里的内部缓冲区分配。这是最隐蔽的坑。插件执行时拿到的是Dims* inputs但很多人不会拿它来分配线程块网格或内部缓冲区而是沿用构造函数里预设的mMaxBatchSize、mBatchSize。旧插件代码里常用的套路是mBatchSize 16然后在enqueue时用16去算显存偏移。这导致即便外部动态batch传了1插件内部依然按16去分配和计算输出是正确了但时间和显存都被白白浪费如果内部逻辑依赖batch做循环结果直接错误。必须改成在enqueue阶段从inputs[0].d[0]实时读取int batchSize inputs[0].d[0]; // 内部循环、缓冲区偏移、kernel启动grid全部用batchSize4. 修复思路与最终代码4.1 重写插件基类并解决维度推导修复方案清晰之后我干脆把插件完整改成IPluginV2DynamicExt实现。这个基类要求实现的方法比旧版多一些核心是getOutputDimensions、supportsFormatCombination、configurePlugin、enqueue。其中configurePlugin里会收到所有输入输出的Dims要注意在这里面不要依赖具体的batch数值只保留格式信息。我的实现过程如下class MyPlugin : public IPluginV2DynamicExt { public: Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override { assert(index 0); Dims out inputs[0]; out.d[2] 640; out.d[3] 640; return out; } bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs) override { return inOut[pos].format TensorFormat::kLINEAR inOut[pos].type DataType::kFLOAT; } void configurePlugin(const DynamicPluginTensorDesc* in, int nbInputs, const DynamicPluginTensorDesc* out, int nbOutputs) override { // 这里不要读in[0].dims.d[0]因为configure阶段拿到的可能是min/opt/max之一的shape // 只记录通道数和空间尺寸 } int enqueue(int batchSize, const void* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) override { // enqueue是运行时阶段batchSize由上下文传入 // 内部所有内存offset都按batchSize计算 return 0; } };这里特别说一句enqueue函数的batchSize参数在动态shape模式下依然是有效的它会由ExecutionContext根据set_binding_shape设置的维度自动传进来。但很多旧代码是从构造函数里存下来的mBatchSize取的这就是我前面说的第三处坑的来源。4.2 启用动态维度绑定并按实际shape执行插件修好之后整个动态链路就畅通了。构建和执行的完整流程我整理成了一个可以直接照抄的版本# 构建阶段 builder trt.Builder(TRT_LOGGER) network builder.create_network(1 int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) parser trt.OnnxParser(network, TRT_LOGGER) with open(model.onnx, rb) as f: parser.parse(f.read()) config builder.create_builder_config() config.max_workspace_size 1 30 profile builder.create_optimization_profile() profile.set_shape(input, min(1, 3, 640, 640), opt(4, 3, 640, 640), max(16, 3, 640, 640)) config.add_optimization_profile(profile) engine_bytes builder.build_serialized_network(network, config) # 执行阶段 engine deserialize_engine(engine_bytes) context engine.create_execution_context() context.set_binding_shape(0, (8, 3, 640, 640)) # 这次跑8张图 bindings pre_allocate_buffers(engine, context) outputs context.execute_v2(bindings) # 返回值是bool绑定地址要提前算好注意context.set_binding_shape必须在execute_v2之前调用而且传入的维度必须落在profile设定的min和max范围内。如果某个维度没设置过执行时TensorRT会直接报“incomplete binding”之类的错误。4.3 多种batch的实测结果与耗时对比修好之后我用同一套engine连续测了batch 1、4、8、16以及中间值2每一档都调用了set_binding_shape再执行。实测证明动态切换完全正常没有再出现维度报错。耗时数据我贴在下面可以直观看到动态batch带来的收益实际batch单batch平均耗时ms总吞吐张/秒显存占用MB16.216181227.8256924411.43501136819.940215721636.54382410在batch1时单张延迟保持在个位数毫秒适合低延迟场景batch提高到16后吞吐翻了接近三倍。线上服务的调度策略后来就做成低峰期用小batch高峰期根据排队长度自动切到更大的batch整机GPU利用率明显改善。5. 留给后来者的检查清单和一句大实话5.1 动态batch排查自检清单把这次排查过程中所有踩过的点收敛成一张清单任何一个环节不对动态batch都会悄悄退化成固定shape。我每次遇到类似问题都会按这个顺序检查检查项具体表现处理方式是否开启EXPLICIT_BATCH不开启时ONNX的N维会被抹去create_network(1 int(flag))profile的min/opt/max是否对齐opt值会被当作默认固定shape优化点应选真实峰值batch附近解析日志里是否有“WARNING”解析器遇到不支持的算子会警告逐个处理警告对应的层插件是否继承DynamicExt旧版插件不参与动态shape推导重写为IPluginV2DynamicExt插件getOutputDimensions是否引用d[0]写死batch会让后续全链固定改为用inputs[0].d[0]传递enqueue里是否使用传入的batchSize使用成员变量会导致内部计算错误所有偏移和网格都按运行时batch算是否对中间层shape做过断言固定shape经常在中间层出现对关键tensor调用getShape校验5.2 为什么TensorRT7的动态batch不是纯配置问题最后说句大实话网上大量教程把动态batch简化成“加profile就好”很大程度上只适用于全标准卷积、没有插件、没有复杂shape变换的简单模型。一旦模型里出现Reshape、自定义插件、动态分块这类结构动态batch就从“配置问题”变成了“网络结构问题”。我后来碰到的很多类似案例根因都出在模型的某个中间张量维度上——要么是Reshape把batch混进了其他维要么是插件解析时把某个维度锁死。问题现象一模一样都是engine绑定shape变固定但定位路径各有不同。最通用的办法就是我这篇文章里用的先造最小网络确认平台能力再二分删减找嫌疑层最后逐个检查插件里的维度传递。这套方法在TensorRT7上管用换到新版TensorRT的OnnxParser时代同样适用。如果你也正在被这个问题折磨建议直接把第三节的排除法流程跑一遍大概率会比我当时的弯路少走一半。
