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

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

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

在深度学习模型的实际应用中，训练只是第一步，如何将模型高效部署到生产环境才是关键。本文将以经典的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)	FPS
pytorch	11.5977, 23.6635, -19.3974, -45.7438, -55.7060	9	167
onnx	11.5977, 23.6634, -19.3974, -45.7438, -55.7060	4.4	227
TensorRT（FP32）	11.5977, 23.6634, -19.3972, -45.7438, -55.7060	0.6	1666
TensorRT（FP16）	10.5000, 23.5313, -20.0625, -45.84375, -55.750	0.2	5000

可以看到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 = "https://storage.googleapis.com/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的静态输入形状不一致