cap2：1000分类的ResNet的TensorRT部署指南（python版）

《TensorRT全流程部署指南》专栏文章目录：

cap1：TensorRT介绍及CUDA环境安装cap2：1000分类的ResNet的TensorRT部署指南（python版）cap3：自定义数据集训练ResNet的TensorRT部署指南（python版）cap4：YoloV5的TensorRT部署指南（python版） …

文章目录 1、保存pytorch模型1.1 获取pth模型1.2 建立标杆 2、导出ONNX2.1 导出模型2.2 验证模型2.3 可视化模型结构 3、环境搭建3.1 TensorRT的安装3.2 安装pycuda 4、转换TensorRT引擎4.1 使用trtexec工具完成序列化4.2 使用python的API进行转换 5、推理5.1 推理代码5.2 结果对比5.3 图片推理 6、多batchsize情况6.1 onnx导出设置6.2 TensorRT序列化设置6.3 推理设置常见问题 在深度学习模型的实际应用中，训练只是第一步，如何将模型高效部署到生产环境才是关键。本文将以经典的ResNet18模型（1000分类）为例，手把手带你走完从模型权重获取、优化到使用TensorRT实现高性能推理的完整流程。无论你是刚接触模型部署的新手，还是希望优化推理性能的开发者，本文都将为你提供清晰、实用的指导，助你快速掌握工业级模型部署的核心技能！

环境：ubuntu20.04+4070Ti

1、保存pytorch模型 1.1 获取pth模型

首先，需要获取一个PyTorch模型。以ResNet18为例，可以直接使用torchvision中提供的预训练模型，该模型默认支持1000分类。以下为示例代码。在实际应用中，可以在ResNet18的基础上进行微调，或者从头训练模型，最终保存为.pth格式的权重文件即可。

默认情况下，ResNet18 的输入是一个 4D 张量，形状为 (batch_size, 3, 224, 224)，输出是输出形状：(batch_size, 1000)。一般使用单batch，所以输入为(1,3,224,224)，输出为(1,1000)。

import torch import torchvision.models as models # 加载预训练的 ResNet18 模型 resnet18 = models.resnet18(pretrained=True) # 保存模型为 .pth 文件 torch.save(resnet18.state_dict(), 'resnet18.pth') print("ResNet18 模型已成功导出为 resnet18.pth") # 建立标杆 resnet18.eval() # 创建一个全部值为 127 的输入张量 (batch_size=1, channels=3, height=224, width=224) input_tensor = torch.full((1, 3, 224, 224), 127.0) # 注意：值为浮点数 127.0 # 运行推理 with torch.no_grad(): # 禁用梯度计算 output = resnet18(input_tensor) print("输出张量形状:", output.shape) # 输出: torch.Size([1, 1000]) print("输出张量部分值:", output[0, :5]) # 打印前 5 个类别的分数 # 11.5977, 23.6635, -19.3974, -45.7438, -55.7060 1.2 建立标杆

在上面的代码中，后半段是建立标杆。这里建立标杆是为了验证后续转换精度。同一个卷积在不同的框架中实现方式可能略有不同，那么模型从这个框架转换到另一个框架，其性能不能保证100%不变，肯定是存在一定的变化的。比如给定图片在pytorch中推理，得分为0.95，在tensorrt部署后可能得分为0.946或者0.96都有可能。造成这种原因是多方面的，最常见的是两种：1）算子本身实现方式不一样；2）图像预处理不一样。比如缩放有很多种算法，使用python的cv2进行缩放，或者使用c++自己实现的缩放，甚至其他库实现的缩放，3种算法的最终矩阵不可能完全一致，轻微的像素值差异都会导致最终结果的变化。

这里使用值全部为127的(1,3,224,224)进行推理，得到结果。在onnx和tensorrt部署后我们也可以使用值全部为127的输入进行推理，看看结果和这里的结果是否一样，如果差距很大说明转换过程中出现了问题，需要及时排查。

2、导出ONNX 2.1 导出模型

ONNX是一个中间件，比如下图只要各深度学习框架支持onnx，那么就能够实现互通，大大方便框架的维护和使用者的部署。像本文就是pytorch–>onnx–>tensorrt实现最终的模型转换和部署。 pytorch提供了onnx转换API，可以轻松转换：

import torch import torchvision.models as models # 1. 加载模型结构，所谓的结构由模型架构+模型参数组成。 resnet18 = models.resnet18(pretrained=False) # 不加载预训练权重 weight_path = "resnet18.pth" # 替换为你的 .pth 文件路径 resnet18.load_state_dict(torch.load(weight_path)) # 2. 将模型设置为评估模式 resnet18.eval() # 3. 创建一个虚拟输入张量 (batch_size=1, channels=3, height=224, width=224) dummy_input = torch.randn(1, 3, 224, 224) # 5. 定义导出的 ONNX 文件名 onnx_filename = "resnet18.onnx" # 6. 导出模型为 ONNX 格式 torch.onnx.export( resnet18, # 要导出的模型 dummy_input, # 虚拟输入 onnx_filename, # 导出的 ONNX 文件名 export_params=True, # 导出模型参数（权重） opset_version=11, # ONNX 算子集版本（推荐使用 11 或更高版本） do_constant_folding=True, # 是否优化常量折叠 input_names=["input"], # 输入节点的名称 output_names=["output"], # 输出节点的名称 ) print(f"模型已成功从 {weight_path} 加载权重并导出为 {onnx_filename}") 2.2 验证模型

导出成功后，我们可以验证一下onnx模型是否正常。可以使用onnxruntime推理引擎简单推理一下：

import onnx import onnxruntime as ort import numpy as np # 1. 用checker方法检查模型是否有效 onnx_model = onnx.load("resnet18.onnx") onnx.checker.check_model(onnx_model) # 检查模型是否有效，如果无效会报错 # 2. 用模拟推理的方式检查模型是否有效 ort_session = ort.InferenceSession("resnet18.onnx") # 3. 准备输入数据 input_name = ort_session.get_inputs()[0].name input_data = np.full((1, 3, 224, 224),127,np.float32) # 4. 运行推理 outputs = ort_session.run(None, {input_name: input_data}) # 5. 打印输出形状 print(outputs[0][0][:10]) # 展示前10个类别的输出:11.597546 23.663443 -19.39737 -45.74385 -55.70603 print(outputs[0].shape) # 输出: (1, 1000)

我们可以看到onnx模拟推理的输出和pytorch的输出几乎一样，至少保证了小数点后3位。

2.3 可视化模型结构

Netron 是神经网络、深度学习和机器学习模型的查看器。Netron 支持 ONNX、TensorFlow Lite、Core ML、Keras、Caffe、Darknet、PyTorch、TensorFlow.js、Safetensors 和 NumPy。在这里我们使用该工具可视化resnet18.onnx模型结构。netron提供了在线版本（也有桌面应用版本），在网页中直接打开模型自动上传后就可以显示了。这里我们可视化resnet18.onnx的结果如图所示，展示了输入输出两个节点的信息。

3、环境搭建

TensorRT在Nvidia显卡上优化并推理，所以必然需要先配置好显卡驱动、CUDA以及cuDNN等。 这部分的教程可以参考：cap1：TensorRT介绍及CUDA环境安装中的4、基础环境配置

3.1 TensorRT的安装

首先去NVIDIA TensorRT Download选择一个版本，比如TensorRT8。我是Ubuntu20，但是选择TAR包，所以选择Linux_X86_64+TAR package下载。 下载并解压后，有多个目录，其中：

bin：有trtexec可执行文件，通过这个工具可以实现序列化转换include和lib：是C++开发用的头文件和库文件python：python的whl文件，通过pip安装

在python中根据python版本安装对应的whl文件，比如我的是python3.10则pip install tensorrt-8.6.1-cp310-none-linux_x86_64.whl

3.2 安装pycuda

TensorRT是在nvidia显卡上推理，要求直接输入输出都在GPU显存空间中。所以会使用CUDA进行相关操作，pycude就是python操作CUDA的库。

安装方式非常简单：pip install pycuda

4、转换TensorRT引擎

这一步主要完成红框内的工作，就是将onnx模型序列化为TensorRT的engine模型。在这个过程中会自动根据设备的特异性进行优化，该过程十分缓慢。但是完成后会保存为本地文件，下次可以直接加载使用，不用再次序列化。 这个过程可以使用基于python的API完成，或者直接使用trtexec工具。

4.1 使用trtexec工具完成序列化

在第3节中下载了TensorRT包，在bin中有trtexec工具。打开终端进入trtexec工具所在文件夹就可以使用该命令工具了。

转换命令为： trtexec --onnx=resnet18.onnx --saveEngine=resnet18.engine --fp16 onnx指定onnx模型路径，saveEngine知道转换后engine模保存路径，fp16表示使用FP16，这个可以大大提高推理速度。不加–fp16则默认使用FLOAT32

转换完成后如图，最后会出现PASSED结果，此时就得到engine模型。

4.2 使用python的API进行转换

使用python的API进行转换具有比较固定的流程，只不过可以根据需求比如动态输入等进行相应的设置。最终结果和trtexec转换没有什么区别。

不管什么方式转换得到的engine都可以被python或者c++拿去部署使用，但是要注意：1）版本对应：即转换使用的TensorRT版本和部署推理使用的版本需要一致；2）不支持跨设备：即在哪种设备上转换的就只能在哪种设备上使用，因为转换时根据设备特异性进行了优化。

import tensorrt as trt # import pycuda.driver as cuda # import pycuda.autoinit # TensorRT 需要一个 日志对象，用于输出 WARNING 级别的日志，帮助调试问题 TRT_LOGGER = trt.Logger(trt.Logger.WARNING) def build_engine(onnx_file_path, engine_file_path): with trt.Builder(TRT_LOGGER) as builder, \ builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) as network, \ trt.OnnxParser(network, TRT_LOGGER) as parser: # config 配置 TensorRT 编译选项 builder_config = builder.create_builder_config() builder_config.set_flag(trt.BuilderFlag.FP16) # 设置FP16加速（如果 GPU 支持），提高计算速度并减少显存占用。注释则默认使用FP32 builder_config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30) # 设置最大工作空间（1GB），用于存储中间 Tensor 和计算资源 # 读取 ONNX 文件 with open(onnx_file_path, 'rb') as model: if not parser.parse(model.read()): for error in range(parser.num_errors): print(parser.get_error(error)) return None # 构建 TensorRT 引擎 serialized_engine = builder.build_serialized_network(network, builder_config) # 开始反序列化为可执行引擎 runtime = trt.Runtime(TRT_LOGGER) engine = runtime.deserialize_cuda_engine(serialized_engine) # 序列化并保存引擎 with open(engine_file_path, 'wb') as f: f.write(engine.serialize()) print(f"Saved TensorRT engine to {engine_file_path}") if __name__ == "__main__": onnx_file = "resnet18.onnx" engine_file = "resnet18.engine" build_engine(onnx_file, engine_file)

通过这种方式，我们也成功得到的engine模型。

5、推理 5.1 推理代码

推理的代码十分简单，推理的流程图如下。TensorRT的推理是在GPU上完成，所以推理只能从GPU的显存获取数据，推理的直接输出也是在显存中。所以输入数据需要从RAM（cpu端，或者host端）复制到GPU（设备端）中去。推理后需要从GPU复制到cpu后，才能使用numpy或者其他库进行操作。

为什么要分配显存空间，因为CUDA程序只能读取位于显存上的数据，而一般程序读取得数据都位于RAM。所以我们需要先在CUDA上用cuda.mem_alloc为模型推理的输入和输出开辟足够的空间，然后将python读取的图像数据用cuda.memcpy_htod_async复制到CUDA显存中去，推理时直接从显存读取数据。推理结束后，输出放在显存上，再使用cuda.memcpy_dtoh_async将结果从显存复制到RAM中。同理接收数据时，需要再RAM中开辟足够的空间接收从显存来的数据，所以使用output = np.empty((1, 25200, 85), dtype=np.float32)开辟了足够大小的空间。

完整的推理代码如下所示：

import tensorrt as trt import pycuda.driver as cuda import pycuda.autoinit # cuda初始化 import numpy as np TRT_LOGGER = trt.Logger(trt.Logger.WARNING) def load_engine(engine_file_path): """ 从engine模型反序列化 :param engine_file_path: :return: """ """加载 TensorRT 引擎""" with open(engine_file_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime: return runtime.deserialize_cuda_engine(f.read()) def allocate_buffers(engine): """ 给模型直接输入和输出分配CUDA缓存区 :param engine: :return: """ inputs, outputs, bindings = [], [], [] stream = cuda.Stream() for binding in engine: # 遍历模型engine中的所有输入输出节点 size = trt.volume(engine.get_tensor_shape(binding)) * np.dtype(np.float32).itemsize # 602112 device_mem = cuda.mem_alloc(size) # 分配size大小的空间，用于储存该节点的数据 engine.get_tensor_mode(binding) if engine.get_tensor_mode(binding) == trt.TensorIOMode.INPUT: inputs.append(device_mem) # 如果是输入节点则将该空间地址存入inputs elif engine.get_tensor_mode(binding) == trt.TensorIOMode.OUTPUT: outputs.append(device_mem) # 如果是输出节点则将该空间地址存入outputs bindings.append(int(device_mem)) return inputs, outputs, bindings, stream def infer(engine, image): """执行 TensorRT 推理""" context = engine.create_execution_context() # 创建推理的上下文环境 inputs, outputs, bindings, stream = allocate_buffers(engine) # 初始化分配CUDA显存空间 # 将输入数据拷贝到 GPU cuda.memcpy_htod_async(inputs[0], image, stream) # image是cpu的ram数据，需要拷贝到GPU，才能被tensorrt推理使用 # 执行推理 context.execute_async_v2(bindings, stream.handle, None) # 推理 # 从 GPU 拷贝输出结果 output = np.empty((1, 1000), dtype=np.float32) # 初始化一个cpu的ram空间，方便从gpu接受输出数据 cuda.memcpy_dtoh_async(output, outputs[0], stream) # 将位于gpu上的输出数据复制到cpu上，方便对齐操作 stream.synchronize() # 同步，gpu是并行的，这里相当于等待同步一下 return output # 最后将输出结果返回 if __name__ == "__main__": engine_path = "resnet18.engine" image_path = "test.jpg" # 替换成你的图片路径 print("Loading TensorRT engine...") engine = load_engine(engine_path) print("Preprocessing image...") input_data = np.full((1, 3, 224, 224),127,np.float32) print("Running inference...") output = infer(engine, input_data) print("Inference completed! Top-5 predictions:") print(output[0][:5])

在实际使用过程中，下面可以在最开始执行一次即可，创建的上下文环境可以重复使用，分配的显存空间也可以重复使用，这样就降低了这两个命令的耗时。

context = engine.create_execution_context() # 创建推理的上下文环境 inputs, outputs, bindings, stream = allocate_buffers(engine) # 初始化分配CUDA显存空间

最终我们可以将infer函数改写为类：

class InferClas: def __init__(self,engine): self.context = engine.create_execution_context() # 创建推理的上下文环境 self.inputs, self.outputs, self.bindings, self.stream = allocate_buffers(engine) # 初始化分配CUDA显存空间 def infer(self,image): # 将输入数据拷贝到 GPU cuda.memcpy_htod_async(self.inputs[0], image, self.stream) # image是cpu的ram数据，需要拷贝到GPU，才能被tensorrt推理使用 # 执行推理 self.context.execute_async_v2(self.bindings, self.stream.handle, None) # 推理 # 从 GPU 拷贝输出结果 output = np.empty((1, 1000), dtype=np.float32) # 初始化一个cpu的ram空间，方便从gpu接受输出数据 cuda.memcpy_dtoh_async(output, self.outputs[0], self.stream) # 将位于gpu上的输出数据复制到cpu上，方便对齐操作 self.stream.synchronize() # 同步，gpu是并行的，这里相当于等待同步一下 return output # 最后将输出结果返回 # 推理的调用 inferclas=InferClas(engine) # 初始化一次 for i in range(10): # 后续的执行都只需要调用infer函数，提高每帧的推理速度 output = inferclas.infer(input_data) 5.2 结果对比

至此，我们完成了python的tensorrt部署。我们来看一下模型的直接输出情况。 # fp16:10.5 23.53125 -20.0625 -45.84375 -55.75 # fp32:11.597782 23.66346 -19.397268 -45.743862 -55.706043 我们还是以13224*224的全部为127的模拟图像作为输入，并测试resnet18在pytorch、onnx、TensorRT（FP32）和TensorRT（FP16）下的直接输出和速度比较。

速度统计需求如下，因为第一次推理存在冷启动过程，耗时增大，故舍弃，只看最后10次的时间均值。设备为i7-12700KF+4070Ti

# pytorch with torch.no_grad(): # 禁用梯度计算 for i in range(11): t1=time.time() output = resnet18(input_tensor) print(time.time()-t1) # onnx for i in range(11): t1=time.time() outputs = ort_session.run(None, {input_name: input_data}) print(time.time()-t1) #TensorRT inferclas=InferClas(engine) for i in range(11): t1=time.time() output = inferclas.infer(input_data) print(time.time()-t1) 模型输出耗时(ms)FPSpytorch11.5977, 23.6635, -19.3974, -45.7438, -55.70609167onnx11.5977, 23.6634, -19.3974, -45.7438, -55.70604.4227TensorRT（FP32）11.5977, 23.6634, -19.3972, -45.7438, -55.70600.61666TensorRT（FP16）10.5000, 23.5313, -20.0625, -45.84375, -55.7500.25000

可以看到TensorRT部署后推理速度有了极大的提高，能够达到上千帧。

但是注意，根据上面的流程图，每一帧的推理还包含其他步骤，所以实际帧率会低。在预处理部分也会影响到帧率，如何更高效率进行预处理也是模型部署的关键一环，此处不再细谈。

5.3 图片推理

图片推理和推理代码一样，只不过需要对图像进行预处理操作，完整代码如下。get_input函数的功能：读取一张图片，缩放到224*224，归一化到0~1，再Normalize，最后改变维度得到（1，3，224，224）。

import time import tensorrt as trt import pycuda.driver as cuda import pycuda.autoinit # cuda初始化 import numpy as np import sys TRT_LOGGER = trt.Logger(trt.Logger.WARNING) def load_engine(engine_file_path): """ 从engine模型反序列化 :param engine_file_path: :return: """ """加载 TensorRT 引擎""" with open(engine_file_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime: return runtime.deserialize_cuda_engine(f.read()) def allocate_buffers(engine): """ 给模型直接输入和输出分配CUDA缓存区 :param engine: :return: """ inputs, outputs, bindings = [], [], [] stream = cuda.Stream() for binding in engine: # 遍历模型engine中的所有输入输出节点 size = trt.volume(engine.get_tensor_shape(binding)) * np.dtype(np.float32).itemsize # 602112 device_mem = cuda.mem_alloc(size) # 分配size大小的空间，用于储存该节点的数据 engine.get_tensor_mode(binding) if engine.get_tensor_mode(binding) == trt.TensorIOMode.INPUT: inputs.append(device_mem) # 如果是输入节点则将该空间地址存入inputs elif engine.get_tensor_mode(binding) == trt.TensorIOMode.OUTPUT: outputs.append(device_mem) # 如果是输出节点则将该空间地址存入outputs bindings.append(int(device_mem)) return inputs, outputs, bindings, stream class InferClas: def __init__(self,engine): self.context = engine.create_execution_context() # 创建推理的上下文环境 self.inputs, self.outputs, self.bindings, self.stream = allocate_buffers(engine) # 初始化分配CUDA显存空间 def infer(self,image): # 将输入数据拷贝到 GPU cuda.memcpy_htod_async(self.inputs[0], image, self.stream) # image是cpu的ram数据，需要拷贝到GPU，才能被tensorrt推理使用 # 执行推理 self.context.execute_async_v2(self.bindings, self.stream.handle, None) # 推理 # 从 GPU 拷贝输出结果 output = np.empty((1, 1000), dtype=np.float32) # 初始化一个cpu的ram空间，方便从gpu接受输出数据 cuda.memcpy_dtoh_async(output, self.outputs[0], self.stream) # 将位于gpu上的输出数据复制到cpu上，方便对齐操作 self.stream.synchronize() # 同步，gpu是并行的，这里相当于等待同步一下 return output # 最后将输出结果返回 def get_input(img_path): """ 读取一张图片，缩放到224*224，归一化到0~1，再Normalize，最后改变维度得到（1，3，224，224） :param img_path: :return: """ import cv2 # 1. 使用 OpenCV 读取图像并缩放到 224x224 img = cv2.imread(img_path) img_resized = cv2.resize(img, (224, 224)) # 2. 使用 NumPy 将图像转为张量格式 (C, H, W) -> (3, 224, 224) # 转换为 [0, 1] 范围内 img_normalized = img_resized.astype(np.float32) / 255.0 # 将像素值从 [0, 255] 转换到 [0, 1] # 3. Normalize 操作：使用 ImageNet 均值和标准差 mean = [0.485, 0.456, 0.406] std = [0.229, 0.224, 0.225] img_normalized = (img_normalized - mean) / std # 归一化 # 4. 改变维度从 (H, W, C) -> (1, C, H, W) img_tensor = np.transpose(img_normalized, (2, 0, 1)) # 转换通道到最前面 img_tensor = np.expand_dims(img_tensor, axis=0) # 添加 batch 维度，变为 (1, 3, 224, 224) # 5. 确保数组是连续的,memcpy_htod_async要求的 img_tensor=img_tensor.astype(np.float32) img_tensor = np.ascontiguousarray(img_tensor) return img_tensor def get_class_name(idx): # 加载ImageNet类别标签 import requests # 下载ImageNet类别标签 labels_url = " storage.googleapis /download.tensorflow.org/data/imagenet_class_index.json" labels_map = requests.get(labels_url).json() # 将类别索引映射到类别名称 labels = [labels_map[str(i)][1] for i in range(len(labels_map))] # 输出预测结果 print(f"Predicted class id: {idx} Predicted class: {labels[idx]}") if __name__ == "__main__": engine_path = "resnet18_fp16.engine" print("Loading TensorRT engine...") engine = load_engine(engine_path) print("Preprocessing image...") input_data = get_input("rabbit2.jpeg") print("Running inference...") inferclas=InferClas(engine) output = inferclas.infer(input_data) predicted_idx=np.argmax(output) get_class_name(predicted_idx) # 从googleapi获取imagenet标签进行转换，可能存在网络不通畅的问题

对于下面的图片，推理结果为：Predicted class id: 332 Predicted class: Angora，即安哥拉兔

6、多batchsize情况

以上介绍的都是输入为（1，3，224，224）的情况，batchsize为1。如果希望推理时batchsize大于1，则需要另外处理：1）onnx环节；2）TensorRT序列化环节；3）推理处理环节；

多batchsize实际上就是动态设置batchsize维度，希望这个维度是可以变动的，而不是固定为1。

6.1 onnx导出设置

onnx导出时，只需要添加dynamic_axes参数，这个参数以字典形式设置输入节点input和输出节点output的0轴（即batchsize轴）为动态轴。后面的“batch_size”只是描述信息。注意这里的输入和输出需要同时设置。

dummy_input = torch.randn(1, 3, 224, 224) # 6. 导出模型为 ONNX 格式 torch.onnx.export( resnet18, # 要导出的模型 dummy_input, # 虚拟输入 onnx_filename, # 导出的 ONNX 文件名 export_params=True, # 导出模型参数（权重） opset_version=11, # ONNX 算子集版本（推荐使用 11 或更高版本） do_constant_folding=True, # 是否优化常量折叠 input_names=["input"], # 输入节点的名称 output_names=["output"], # 输出节点的名称 dynamic_axes={ # 动态轴（支持动态 batch_size） "input": {0: "batch_size"}, "output": {0: "batch_size"}, }, ) 6.2 TensorRT序列化设置

TensorRT序列化设置也非常简单，实际上就是在builder_config中添加一个profile。profile也是一个序列化的配置参数，功能就是为输入节点设置形状信息。

# profile配置 用于配置动态形状输入的优化参数。是config配置的一个子项，最终需要添加到config中 profile = builder.create_optimization_profile() profile.set_shape("input", (10, 3, 224, 224), # 设置最小 shape (10, 3, 224, 224), # 设置最优 shape (10, 3, 224, 224)) # 设置最大 shape builder_config.add_optimization_profile(profile)

profile.set_shape()有四个参数，第一个参数是节点名字，你是为哪个节点设置形状（这里的input在onnx导出时参数input_names所设置的）。第二个参数是这个节点接受的最小形状。第三个参数是这个节点接受的最佳形状。第四个参数是这个节点接受的最大形状。

从这里可以看出，如果我设置profile.set_shape("foo", (2, 3, 100, 200), (2, 3, 150, 250), (2, 3, 200, 300))表明，foo这个节点固定了batchsize=3，但是宽度允许在100~200，高度允许在200~300。注意相应的在onnx导出时就需要在dynamic_axes将宽和高设置为动态轴。

在文本，如果已知batchsize一定为10，输入形状也不变就是(3,224,224)，故最小最优最大全部设置为期望的形状。

添加了这个profile，再正常序列化即可。

6.3 推理设置

推理设置就涉及到2个地方。

第一，需要一个(10,3,224,224)的输入。可以将10帧图片预处理后堆叠为(10,3,224,224)的形状作为输入。这里使用np模拟：

input_data = np.full((10, 3, 224, 224),127,np.float32) # input_data.shape=(10,3,224,224)

第二，修改输出拷贝

class InferClas: def __init__(self,engine): self.context = engine.create_execution_context() self.inputs, self.outputs, self.bindings, self.stream = allocate_buffers(engine) def infer(self,image): cuda.memcpy_htod_async(self.inputs[0], image, self.stream) self.context.execute_async_v2(self.bindings, self.stream.handle, None) # 之前batchsize=1，所以是np.empty((1, 1000)）,但是现在batchsize=10，就需要np.empty((10, 1000)的空间来接收输出。 # self.outputs[0]中数据量是10*1000，如果output只开辟了1*1000的空间，则最终只有输出的前1000个数据，是不完整的 output = np.empty((10, 1000), dtype=np.float32) # 初始化一个cpu的ram空间，方便从gpu接受输出数据 cuda.memcpy_dtoh_async(output, self.outputs[0], self.stream) self.stream.synchronize() # ...... input_data = np.full((10, 3, 224, 224),127,np.float32) inferclas=InferClas(engine) output = inferclas.infer(input_data) print(output.shape) # (10,1000)

至此，多batchsize的需求处理完毕，其他如宽高等动态尺寸需求同理。

常见问题

[network.cpp::operator()::3212] Error Code 4: Internal Error (input: kMIN dimensions in profile 0 are [10,3,224,224] but input has static dimensions [1,3,224,224].)：onnx导出时没有设置动态轴，在TensorRT序列化时设置的输入形状和onnx的静态输入形状不一致