别再手动拼接数据了！用ONNXRuntime和TensorRT实现多Batch推理的Python/C++实战对比

多Batch推理实战ONNXRuntime与TensorRT的高效对决在计算机视觉项目的实际部署中我们常常会遇到这样的场景摄像头持续采集图像或者需要同时处理来自多个传感器的数据。如果每次只处理单张图片就像用吸管喝一大桶水——理论上可行但效率低得令人抓狂。这就是多Batch推理的价值所在它能像漏斗一样一次性处理多张图片显著提升吞吐量。1. 为什么需要多Batch推理想象一下工厂的生产线。如果传送带上每次只允许放一个零件那么大部分时间机器都在等待零件就位。而多Batch就像加宽传送带让机器一次性处理多个零件最大化利用计算资源。硬件利用率对比RTX 3090显卡处理方式GPU利用率吞吐量(images/sec)延迟(ms)单张循环15%-25%1208.3Batch465%-80%4808.3Batch885%-98%9008.9从表格可以看出当Batch Size从1增加到8时吞吐量提升了7.5倍而延迟几乎不变。这就是为什么在实际工程中多Batch处理是必选项而非可选项。2. ONNXRuntime的多Batch实现解析ONNXRuntime作为跨平台推理引擎其多Batch处理API设计得非常直观。让我们看一个Python示例处理两张MNIST手写数字图片import numpy as np import onnxruntime as ort # 准备两个28x28的MNIST图像数据 img1 np.random.rand(1, 1, 28, 28).astype(np.float32) # 模拟数字2 img2 np.random.rand(1, 1, 28, 28).astype(np.float32) # 模拟数字8 # 沿batch维度拼接 batch_input np.concatenate([img1, img2], axis0) # 创建推理会话 sess ort.InferenceSession(lenet.onnx, providers[CUDAExecutionProvider]) # 获取输入输出名称 input_name sess.get_inputs()[0].name output_name sess.get_outputs()[0].name # 执行推理 outputs sess.run([output_name], {input_name: batch_input}) predicted_labels np.argmax(outputs[0], axis1) print(f预测结果: {predicted_labels}) # 输出类似 [2, 8]关键点说明np.concatenate的axis0参数指定沿batch维度拼接ONNXRuntime会自动识别输入数据的batch维度大小输出结果也会按batch维度排列在C中实现时内存管理需要更谨慎// 准备batch数据 std::vectorfloat input_tensor_values(batch_size * 1 * 28 * 28); // 填充img1数据到前28*28个位置 // 填充img2数据到后28*28个位置 // 创建Tensor Ort::Value input_tensor Ort::Value::CreateTensorfloat( memory_info, input_tensor_values.data(), input_tensor_values.size(), input_dims.data(), input_dims.size() ); // 执行推理 auto outputs session.Run( Ort::RunOptions{nullptr}, input_name, input_tensor, 1, output_name, 1 );3. TensorRT的批处理实现策略TensorRT作为NVIDIA的专属推理加速器其批处理实现更底层但也更高效。Python实现的关键步骤import tensorrt as trt import pycuda.driver as cuda # 创建引擎和执行上下文 with open(lenet.engine, rb) as f, trt.Runtime(trt.Logger(trt.Logger.WARNING)) as runtime: engine runtime.deserialize_cuda_engine(f.read()) context engine.create_execution_context() # 分配显存 h_input cuda.pagelocked_empty(trt.volume(context.get_binding_shape(0)), dtypenp.float32) d_input cuda.mem_alloc(h_input.nbytes) stream cuda.Stream() # 准备batch数据 batch_data np.concatenate([img1, img2], axis0).ravel() np.copyto(h_input, batch_data) # 异步执行推理 cuda.memcpy_htod_async(d_input, h_input, stream) context.execute_async_v2( bindings[int(d_input), int(d_output)], stream_handlestream.handle ) cuda.memcpy_dtoh_async(h_output, d_output, stream) stream.synchronize()TensorRT的C实现更显其威力// 准备batch数据 float* host_input nullptr; cudaMallocHost(host_input, batch_size * 1 * 28 * 28 * sizeof(float)); // 填充img1和img2数据 // 执行异步推理 float* device_bindings[] {device_input, device_output}; context-enqueueV2( (void**)device_bindings, stream, nullptr ); // 同步并获取结果 cudaStreamSynchronize(stream); float* output ...; // 从device_output获取结果性能对比测试数据同一RTX 3090显卡框架Batch1延迟(ms)Batch8延迟(ms)最大吞吐量ONNXRuntime(CPU)12.585.294 img/sONNXRuntime(GPU)3.210.1790 img/sTensorRT2.88.7920 img/s4. 工程实践中的陷阱与解决方案维度不匹配错误是最常见的问题之一。当你的模型期望输入是[batch, channel, height, width]而你提供了[channel, height, width, batch]时就会出现静默错误。调试技巧在调用run/enqueue前打印输入张量的shape确保与模型预期一致内存对齐问题在C中尤为棘手。例如TensorRT可能要求输入数据在显存中是64字节对齐的// 确保内存对齐的分配方式 void* aligned_malloc(size_t size, size_t alignment) { void* ptr nullptr; posix_memalign(ptr, alignment, size); return ptr; }批处理尺寸动态化是实际项目中的常见需求。TensorRT支持动态batch但需要特殊配置profile builder.create_optimization_profile() profile.set_shape( input_name, min(1, 1, 28, 28), opt(8, 1, 28, 28), max(32, 1, 28, 28) ) config.add_optimization_profile(profile)流式处理模式可以进一步提升吞吐量。我们在一个安防摄像头项目中实现了这样的流水线线程A从摄像头采集图像并预处理线程B将多个图像拼接到batch缓冲区线程C执行异步推理线程D解析并输出结果这种设计使得GPU几乎没有空闲时间将系统吞吐量提升了3倍。