借助TensorRT优化模型推理性能_容器服务 Kubernetes 版 ACK(ACK)-阿里云帮助中心

TensorRT优化模型过程，首先将PyTorch（或TensorFlow）等训练框架训练完成后的模型编译为TensorRT的格式，然后利用TensorRT推理引擎运行这个模型，从而提升这个模型在英伟达GPU上运行的速度，适用于对实时性要求较高的场景。那么该如何借助TensorRT优化模型推理性能呢？本文将演示模型训练编译过程，然后介绍一些TensorRT常用的模型推理性能优化建议。

阅读前提示

为了帮助您更好地理解TensorRT优化模型推理功能，建议您提前了解以下信息：

阅读本文前期望您了解TensorRT简介和TensorRT Cookbook源码等相关内容，以便您对TensorRT的架构和用法有一定的理解。
安装TensorRT还需注意CUDA的兼容性，由于TensorRT是专为英伟达GPU设计的，因此它只能在英伟达的硬件上使用，更多信息可查看TensorRT官方文档。
本文中使用的Nsight Systems软件，主要用于观察全局的Profiling，如核函数读写情况，核函数之间的调度情况，SM占有率，CPU和GPU之间的异步执行的情况等。

本文中加速效果取决于模型的类型和大小，也取决于我们所使用的显卡类型。

模型编译示例

示例拉取现有ResNet18模型进行简单的训练作为演示。您可以使用ResNet18模型跟随示例实现模型性能分析与优化的思路和技巧。
- TensorRT版本为v8.6.1，更多版本请参考TensorRT下载。
- PyTorch版本为2.2.0。
- GPU卡型号为V100-SXM2-32GB，并使用英伟达官方PyTorch镜像运行代码。
  Docker拉取英伟达官方PyTorch镜像：docker pull nvcr.io/英伟达/PyTorch:24.01-py3。注意启动Docker时，需为容器挂载Shm（docker run --shm-size=）和共享宿主机IPC（--ipc=host）。

训练模型并生成ONNX格式模型文件。

下面代码演示了如何拉取现有Resnet18模型并进行简单的训练，最后将模型保存为ONNX格式。

展开查看示例代码

import torch
import torch.nn
import torch.optim
import torch.utils.data
import torchvision.datasets
import torchvision.models
import torchvision.transforms as T

'''
定义一个变换（transform）序列，包含三个步骤：
  (1) T.Resize(224)：将输入图像调整为224x224像素。
  (2) T.ToTensor()：将图像数据转换为PyTorch张量，并且把像素值从[0, 255]范围线性归一化到[0, 1]。
  (3) T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))：标准化图像张量，使其具有指定的均值(0.5, 0.5, 0.5)和标准差(0.5, 0.5, 0.5)，即进行中心化和缩放
'''
transform = T.Compose(
    [T.Resize(224),
     T.ToTensor(),
     T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
# 加载CIFAR10训练数据集，同时下载数据集到本地./data目录，并应用上述定义的图像变换。
train_set = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
# 创建一个DataLoader对象，用于批量加载训练数据，每批大小为32，并且在每个epoch开始时对数据进行随机打乱。
train_loader = torch.utils.data.DataLoader(train_set, batch_size=32, shuffle=True)

# 定义设备变量，指定使用第一个CUDA兼容的GPU进行训练。
device = torch.device("cuda:0")
# 加载一个预训练的ResNet-18模型，并将模型的参数转移到指定的GPU。
model = torchvision.models.resnet18(weights='IMAGENET1K_V1').cuda(device)
# 创建一个交叉熵损失函数对象，用于训练分类任务，并将其转移到GPU。
criterion = torch.nn.CrossEntropyLoss().cuda(device)
# 创建一个随机梯度下降优化器，设置学习率为0.001，动量为0.9，用于优化模型的参数。
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
# 将模型设置为训练模式，这将启用dropout等层的训练行为。
model.train()

# 定义一个训练函数train，它接收一批数据作为参数。
def train(data):
    # 从传入的数据中提取输入和标签，并将它们转移到GPU。
    inputs, labels = data[0].to(device=device), data[1].to(device=device)
    # 将输入数据通过模型进行前向传播，得到输出结果。
    outputs = model(inputs)
    # 用损失函数计算模型输出和真实标签之间的损失值。
    loss = criterion(outputs, labels)
    # 在计算新的梯度之前，先将之前所有参数的梯度清零。
    optimizer.zero_grad()
    # 反向传播，计算损失相对于模型参数的梯度。
    loss.backward()
    # 根据计算得到的梯度更新模型的参数。
    optimizer.step()

# 对每个批次调用train函数进行训练。
for step, batch_data in enumerate(train_loader):
    train(batch_data)

# 将训练好的PyTorch模型导出至ONNX（Open Neural Network Exchange）格式，这种格式可以用来在不同深度学习框架之间交换模型。
# 创建一个随机生成的张量dummy_input，该张量的形状为(1, 3, 224, 224)，这代表了一个具有3个颜色通道(C)，224x224像素的单个图像(batch size为1)。
# 这个张量被用作模型导出时的示例输入，它帮助ONNX确定输入的形状和布局。这个张量随后被转移到之前定义的设备上(device)，在这个案例中是GPU。
dummy_input = torch.randn(1, 3, 224, 224).to(device)
# 定义了一个包含单个字符串元素"input0"的列表input_names，
# 这个字符串将用作导出的ONNX模型中输入张量的名称。
input_names = [ "input0" ]
# 定义了一个包含单个字符串元素"output0"的列表output_names，
# 它将用作ONNX模型中输出张量的名称。
output_names = [ "output0" ]

'''
这是调用PyTorch的ONNX.export函数，用于将PyTorch模型导出为ONNX格式的模型。参数解释如下：
  (1) model：要导出的PyTorch模型，前面代码块中定义并训练的模型。
  (2) dummy_input：模型的示例输入数据。
  (3) 'resnet18.ONNX'：导出的ONNX模型的文件名。
  (4)	verbose=True：如果设置为True，将会打印出在转换过程中的详细信息。
  (5) input_names=input_names：为ONNX模型中的输入指定名称。
  (6) output_names=output_names：为ONNX模型的输出指定名称。
  (7) dynamic_axes={'input0': {0: "nBatchSize"}}：指定模型的输入张量'input0'在其第0维上是动态的，这允许模型处理不同的批量大小。"nBatchSize"是这个动态轴的名称。
'''
torch.ONNX.export(model, dummy_input, 'resnet18.ONNX', verbose=True, input_names=input_names, output_names=output_names,dynamic_axes={'input0': {0: "nBatchSize"}})

保存TensorRT编译模型代码。

编译模型代码如下，并保存在0_build.py文件中。

展开查看示例代码

import argparse  # 导入用于解析命令行参数的库
import os  # 导入用于文件和目录操作的库
import tensorrt as trt  # 导入TensorRT库

# 定义构建函数用于构建TensorRT引擎
def build(logger, ONNX_file, min_shape, optim_shape, max_shape, num_aux_stream, share_profile, fp16):
    errors = []  # 初始化错误列表
    builder = trt.Builder(logger)  # 创建TensorRT Builder对象
    # 创建网络定义，显式指定batch大小
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    profile = builder.create_optimization_profile()  # 创建优化配置文件
    config = builder.create_builder_config()  # 创建构建配置对象

    # 如果启用共享优化配置文件，则在构建配置中设置相应标志
    if share_profile:
        print("enable share profile")
        config.set_preview_feature(trt.PreviewFeature.PROFILE_SHARING_0806, True)

    # 如果设置了辅助流数量（用于流并行执行），则在配置中进行设置
    if num_aux_stream > 0:
        print("set aux stream " + str(num_aux_stream))
        config.max_aux_streams = num_aux_stream

    # 如果启用FP16模式，为构建配置设置FP16标志
    if fp16:
        config.set_flag(trt.BuilderFlag.FP16)

    # 创建ONNX解析器并将其与网络定义和日志记录器关联
    parser = trt.ONNXParser(network, logger)
    
    # 检查ONNX文件是否存在
    if not os.path.exists(ONNX_file):
        errors.append("Failed to find ONNX File!")
        return None, errors

    # 打开ONNX文件，并使用解析器解析文件内容
    with open(ONNX_file, "rb") as model:
        if not parser.parse(model.read()):
            errors.append("failed to parse .ONNX file: ")
            for error in range(parser.num_errors):
                errors.append(parser.get_error(error))
            return None, errors
    
    # 获取网络输入张量并为其设置优化配置文件中的形状
    inputTensor = network.get_input(0)
    profile.set_shape(inputTensor.name, min_shape, optim_shape, max_shape)
    config.add_optimization_profile(profile)  # 将优化配置文件添加到构建配置中
    # 设置详细的性能分析等级
    config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED
    
    # 构建并序列化优化后的TensorRT网络
    engine_string = builder.build_serialized_network(network, config)
    if engine_string == None:
        errors.append("Failed to build engine")
        return None, errors
    
    return engine_string, errors  # 返回构建的引擎字符串和错误列表

# 定义函数用于将引擎字符串保存到文件中
def save_engine(engine_string, planFile):
    with open(planFile, "wb") as f:
        f.write(engine_string)
    return 0

# 定义主函数
def main():
    # 初始化命令行参数解析器
    parser = argparse.ArgumentParser(description='ResNet18 TensorRT Builder')
    # 添加命令行参数
    parser.add_argument('--aux-stream', type=int, default=0, metavar='N',
                        help='specify the aux stream (default: 0)')
    parser.add_argument('--share-profile', action='store_true', default=False,
                        help='enable share profile')
    parser.add_argument('--fp16', action='store_true', default=False,
                        help='enable fp16 mode')
    parser.add_argument('--output', type=str, default='resnet18.plan', metavar='N',
                        help='specify the plan file')
    parser.add_argument('--ONNX-file', type=str, default='resnet18.ONNX', metavar='N',
                        help='specify the ONNX file')
    args = parser.parse_args()  # 解析命令行参数

    logger = trt.Logger(trt.Logger.ERROR)  # 创建TensorRT日志记录器
    # 调用构建函数，传入参数，构建TensorRT引擎
    # 指定profile最小shape：[1, 3, 224, 224]
    # 指定profile最合适shape：[128, 3, 224, 224]
    # 指定profile最大shape：[256, 3, 224, 224]
    engine_string, errors = build(logger, args.ONNX_file, [1, 3, 224, 224], [128, 3, 224, 224], [256, 3, 224, 224], args.aux_stream, args.share_profile, args.fp16)
    
    # 如果构建过程中发生错误，打印错误并退出
    if len(errors) != 0:
        print(errors)
        return 1
    
    save_engine(engine_string, args.output)  # 保存构建的引擎到文件
    return 0  # 返回0，表示成功

# 如果脚本作为主程序运行，则执行主函数
if __name__ == "__main__":
    main()

重要

需要注意的是，由于我们在输出ONNX模型文件时，指定了只有Batch Size可以动态改变，图片通道数（本例为3）、图片宽度和高度（224 x 224）都是固定形状，所以在指定Profile最小（[1, 3, 224, 224]）、最合适（[128, 3, 224, 224]）、最大（[256, 3, 224, 224]）shape时，只有Batch Size做了改变。

保存Baseline模型代码。

Baseline代码如下，保存在1_baseline.py中。

展开查看示例代码

import nvtx  # 导入英伟达 Tools Extension库，用于GPU性能分析
import numpy as np  # 导入NumPy库，用于处理数组和矩阵
import tensorrt as trt  # 导入TensorRT库
from cuda import cudart  # 从cuda模块导入cudart，即CUDA Runtime API

np.random.seed(10088)  # 设置NumPy的随机数种子，确保每次运行生成相同的随机数


# 定义了softmax函数，它用于执行softmax运算
def softmax(x, axis=1):
    e_x = np.exp(x - np.max(x, axis=axis, keepdims=True))   # 为了数值稳定性，减去每个样本的最大值
    return e_x / np.sum(e_x, axis=axis, keepdims=True)  # 计算并返回softmax结果


# 定义用于生成随机输入数据的函数
def data_generation(shape, batches):
    data = []  # 初始化数据列表
    for i in range(batches):  # 循环生成指定批次数量的数据
        # 生成形状为shape的随机数据并转换为float32类型
        data.append(np.random.randn(*shape).astype(np.float32))
    return data  # 返回生成的数据列表


# 定义从磁盘加载预构建的TensorRT推理引擎的函数
def load_engine(logger, plan_file):
    with open(plan_file, "rb") as plan:  # 以二进制读取模式打开引擎文件
        # 使用TensorRT的Runtime，反序列化引擎
        engine = trt.Runtime(logger).deserialize_cuda_engine(plan.read())
    return engine  # 返回加载的引擎


# 定义获取引擎的输入输出张量信息的函数
def get_io_tensors(engine):
    num_io_tensors = engine.num_io_tensors  # 获取引擎的输入输出张量数量
    # 获取所有输入输出张量的名称
    io_tensor_names = [engine.get_tensor_name(i) for i in range(num_io_tensors)]
    # 计算输入张量的数量
    num_input_io_tensors = [engine.get_tensor_mode(io_tensor_names[i]) for i in range(num_io_tensors)].count(trt.TensorIOMode.INPUT)
    return num_io_tensors, io_tensor_names, num_input_io_tensors  # 返回张量信息


# 定义进行推理操作的函数
def infer(engine, data):
    context = engine.create_execution_context()  # 创建TensorRT的执行上下文
    tet = None  # 初始化总时间测量变量
    for i in range(len(data)):  # 遍历每批数据
        if i == 7:  # 特定批次开始时
            # 开始记录NVTX范围（用于性能分析）
            tet = nvtx.start_range(message="Total Elapsed Time(3 batchs)", color="orange")
        nvtx.push_range(message="infer", color="purple")  # 开始记录推理操作的NVTX范围
        infer_once(engine, context, data[i])  # 调用执行单次推理的函数
        nvtx.pop_range()  # 结束NVTX范围记录
    nvtx.end_range(tet)  # 结束总时间的NVTX范围记录


# 定义执行单次推理操作的函数
def infer_once(engine, context, data):
    # 获取引擎的输入输出张量信息
    num_io_tensors, io_tensor_names, num_input_io_tensors = get_io_tensors(engine)
    # 设置输入张量的形状
    context.set_input_shape(io_tensor_names[0], data.shape)
    bufferH, bufferD = [], []  # 初始化主机和设备缓冲区列表
    bufferH.append(data)  # 将输入数据添加到主机缓冲区列表
    # 为输出张量分配空间并添加到主机缓冲区列表
    for i in range(num_input_io_tensors, num_io_tensors):
        bufferH.append(np.empty(context.get_tensor_shape(io_tensor_names[i]), dtype=trt.nptype(engine.get_tensor_dtype(io_tensor_names[i]))))
    # 为每个IO张量在设备上分配内存
    for i in range(num_io_tensors):
        bufferD.append(cudart.cudaMalloc(bufferH[i].nbytes)[1])
    # 将输入数据拷贝到设备缓冲区
    for i in range(num_input_io_tensors):
        cudart.cudaMemcpy(bufferD[i], bufferH[i].ctypes.data, bufferH[i].nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice)
    # 设置执行上下文的张量地址
    for i in range(num_io_tensors):
        context.set_tensor_address(io_tensor_names[i], int(bufferD[i]))
    context.execute_async_v3(0)  # 异步执行推理
    # 将推理结果从设备缓冲区拷贝回主机缓冲区
    for i in range(num_input_io_tensors, num_io_tensors):
        cudart.cudaMemcpy(bufferH[i].ctypes.data, bufferD[i], bufferH[i].nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost)
    cudart.cudaStreamSynchronize(0)  # 同步CUDA流，确保所有操作完成

    # 对输出进行softmax处理并打印结果
    for i in range(num_input_io_tensors, num_io_tensors):
        softmax_scores = softmax(bufferH[i])  # 计算softmax分数
        # 使用np.argmax函数获取概率最高的类别
        predicted_classes = np.argmax(softmax_scores, axis=1)
        max_probs_np = np.max(softmax_scores, axis=1)  # 获取概率最大值
        print("Output Tensor Name: ", io_tensor_names[i])  # 打印输出张量名称
        print("Maximum probability for each image in the batch:\n", max_probs_np)  # 打印每张图像最大概率
        print("Index of predicted class for each image in the batch:\n", predicted_classes)  # 打印每张图像预测的类别索引
    # 释放设备缓冲区的内存
    for b in bufferD:
        cudart.cudaFree(b)
    print("Succeeded running model in TensorRT!")  # 打印成功信息


# 定义主函数，用于生成数据，加载引擎，并执行推理
def main():
    data = data_generation([128, 3, 224, 224], 10)  # 生成数据，10个Batch，每个Batch有128张图片，图片大小为224x224
    logger = trt.Logger(trt.Logger.ERROR)  # 创建TensorRT日志记录器
    engine = load_engine(logger, "resnet18.plan")  # 加载TensorRT引擎
    infer(engine, data)  # 执行推理操作


# 检查脚本是否作为主程序运行并执行主函数
if __name__ == "__main__":
    main()

说明

对于上述代码，有几点需要说明：

利用data_generation随机生成10个Batch，每个Batch有128张图片，图片尺寸为224x224。
我们只观察最后3个Batch，前面的Batch作为Warm Up和Wait使用。

执行推理优化操作并查看过程。
准备完成后，运行如下Shell代码。
```
python 0_build.py 

mkdir -pv reports

nsys profile -w true \
	-t cuda,nvtx,osrt,cudnn,cublas \
	--cuda-memory-usage=true \
	--cudabacktrace=all \
	--cuda-graph-trace=node \
	--gpu-metrics-device=all \
	-f true \
	-o reports/1_baseline \
	python 1_baseline.py
```
运行完成后，在./reports目录下，生成一个名称为1_baseline.nsys-rep文件，可导入Nsight Systems中，Timeline如下。
从上图中可以看到。
- 最后三个Batch总共花费时间约为133.577ms。
- Batch与Batch之间，有一部分时间GPU处于空闲状态（图中标号4的部分），这是由Batch数据传输和Host端打印结果导致的。

模型优化方向

方向1: 重用已分配GPU内存

问题分析。
在Baseline代码中，每次Batch计算都需要重新申请内存，Batch处理完成后，都需要释放数据，GPU内存的申请和释放都是一个比较耗时的操作。

方案设计。

如果能够重用已分配GPU内存，将有利于缩短Batch处理时间。修改Baseline代码，在处理第一个Batch时申请GPU内存，之后的Batch处理所用到的GPU内存都将重用这部分GPU内存。

完整代码如下（注意：只对infer和infer_once做了改动，其他部分代码与Baseline一致），保存在2_reuse_buffers.py中。

展开查看完整代码

import nvtx  # 导入英伟达 Tools Extension库，用于GPU性能分析
import numpy as np  # 导入NumPy库，用于处理数组和矩阵
import tensorrt as trt  # 导入TensorRT库
from cuda import cudart  # 从cuda模块导入cudart，即CUDA Runtime API

np.random.seed(10088)  # 设置NumPy的随机数种子，确保每次运行生成相同的随机数


# 定义了softmax函数，它用于执行softmax运算
def softmax(x, axis=1):
    e_x = np.exp(x - np.max(x, axis=axis, keepdims=True))  # 为了数值稳定性，减去每个样本的最大值
    return e_x / np.sum(e_x, axis=axis, keepdims=True)  # 计算并返回softmax结果


# 定义用于生成随机输入数据的函数
def data_generation(shape, batches):
    data = []  # 初始化数据列表
    for i in range(batches):  # 循环生成指定批次数量的数据
        # 生成形状为shape的随机数据并转换为float32类型
        data.append(np.random.randn(*shape).astype(np.float32))
    return data  # 返回生成的数据列表


# 定义从磁盘加载预构建的TensorRT推理引擎的函数
def load_engine(logger, plan_file):
    with open(plan_file, "rb") as plan:  # 以二进制读取模式打开引擎文件
        # 使用TensorRT的Runtime，反序列化引擎
        engine = trt.Runtime(logger).deserialize_cuda_engine(plan.read())
    return engine  # 返回加载的引擎


# 定义获取引擎的输入输出张量信息的函数
def get_io_tensors(engine):
    num_io_tensors = engine.num_io_tensors  # 获取引擎的输入输出张量数量
    # 获取所有输入输出张量的名称
    io_tensor_names = [engine.get_tensor_name(i) for i in range(num_io_tensors)]
    # 计算输入张量的数量
    num_input_io_tensors = [engine.get_tensor_mode(io_tensor_names[i]) for i in range(num_io_tensors)].count(trt.TensorIOMode.INPUT)
    return num_io_tensors, io_tensor_names, num_input_io_tensors  # 返回张量信息


# 定义进行推理操作的函数
def infer(engine, data):
    context = engine.create_execution_context()  # 创建TensorRT推理上下文
    # 获取引擎的输入输出张量信息
    num_io_tensors, io_tensor_names, num_input_io_tensors = get_io_tensors(engine)
    # 设置第一个输入张量的形状为第一批数据的形状
    context.set_input_shape(io_tensor_names[0], data[0].shape)
    bufferH, bufferD = [], []  # 初始化主机和设备缓冲区列表
    bufferH.append(data[0])  # 在主机缓冲区列表中添加第一批数据
    # 为引擎的其他张量（输出张量）在主机上分配空间
    for i in range(num_input_io_tensors, num_io_tensors):
        bufferH.append(np.empty(context.get_tensor_shape(io_tensor_names[i]), dtype=trt.nptype(engine.get_tensor_dtype(io_tensor_names[i]))))
    # 为所有张量在设备上分配内存
    for i in range(num_io_tensors):
        bufferD.append(cudart.cudaMalloc(bufferH[i].nbytes)[1])
    # 设置执行上下文中每个张量的地址
    for i in range(num_io_tensors):
        context.set_tensor_address(io_tensor_names[i], int(bufferD[i]))
    tet = None  # 初始化一个变量来跟踪NVTX范围
    for i in range(len(data)):  # 遍历所有批次的数据
        if i == 7:  # 当处理到第7批数据时
            # 使用NVTX开始一个范围，用于测量3个批次的总时间
            tet = nvtx.start_range(message="Total Elapsed Time(3 batchs)", color="orange")
        nvtx.push_range(message="infer", color="purple")  # 开始一个NVTX范围来测量单个推理操作
        infer_once(engine, context, bufferH, bufferD, data[i])  # 调用infer_once函数执行单次推理
        nvtx.pop_range()  # 结束NVTX范围
    nvtx.end_range(tet)  # 结束测量3个批次的总时间的NVTX范围
    for b in bufferD:  # 遍历设备缓冲区列表
        cudart.cudaFree(b)  # 释放设备内存


# 定义infer_once函数，它执行一次推理
def infer_once(engine, context, bufferH, bufferD, data):
    # 获取引擎的输入输出张量信息
    num_io_tensors, io_tensor_names, num_input_io_tensors = get_io_tensors(engine)
    bufferH[0] = data  # 更新主机缓冲区列表的第一个元素为当前批次的数据
    # 将输入数据从主机内存拷贝到设备内存
    for i in range(num_input_io_tensors):
        cudart.cudaMemcpy(bufferD[i], bufferH[i].ctypes.data, bufferH[i].nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice)
    context.execute_async_v3(0)  # 在默认CUDA流上异步执行推理
    # 将推理得到的输出数据从设备内存拷贝回主机内存
    for i in range(num_input_io_tensors, num_io_tensors):
        cudart.cudaMemcpy(bufferH[i].ctypes.data, bufferD[i], bufferH[i].nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost)
    cudart.cudaStreamSynchronize(0)  # 同步CUDA流，确保所有之前的操作完成
    nvtx.push_range(message="Print Result", color="green")  # 开始一个NVTX范围来测量结果输出的时间
    print_result(io_tensor_names, num_input_io_tensors, num_io_tensors, bufferH)  # 调用print_result函数输出结果
    nvtx.pop_range()  # 结束NVTX范围
    print("Succeeded running model in TensorRT!")  # 打印成功执行模型的信息


# 定义print_result函数，它输出推理结果
def print_result(io_tensor_names, num_input_io_tensors, num_io_tensors, bufferH):
    for i in range(num_input_io_tensors, num_io_tensors):  # 遍历输出张量
        softmax_scores = softmax(bufferH[i])  # 对输出数据执行softmax运算
        predicted_classes = np.argmax(softmax_scores, axis=1)  # 获取每个样本的预测类别索引
        max_probs_np = np.max(softmax_scores, axis=1)  # 获取每个样本的最高概率
        print("Output Tensor Name: ", io_tensor_names[i])  # 打印输出张量的名称
        print("Maximum probability for each image in the batch:\n", max_probs_np)  # 打印每个样本的最高概率
        print("Index of predicted class for each image in the batch:\n", predicted_classes)  # 打印每个样本的预测类别索引


# 定义main函数，它是程序的入口点
def main():
    data = data_generation([128, 3, 224, 224], 10)  # 生成10批，每批128张3通道224x224大小的随机数据
    logger = trt.Logger(trt.Logger.ERROR)  # 创建TensorRT日志对象，仅记录错误级别的日志
    engine = load_engine(logger, "resnet18.plan")  # 加载序列化的TensorRT引擎
    infer(engine, data)  # 调用infer函数执行推理


if __name__ == "__main__":
    main()

上述代码中，将打印结果放在一个函数当中（print_result）。

准备完成后，运行如下Shell代码。

nsys profile -w true \
	-t cuda,nvtx,osrt,cudnn,cublas \
	--cuda-memory-usage=true \
	--cudabacktrace=all \
	--cuda-graph-trace=node \
	--gpu-metrics-device=all \
	-f true \
	-o reports/2_reuse-buffers \
	python 2_reuse-buffers.py

运行完成后，在./reports目录下，生成一个名称为2_reuse-buffers.nsys-rep文件，导入Nsight Systems中，Timeline如下。从上图中可以看到：最后三个Batch总共花费时间约为128.196ms，比Baseline中减少133.577ms - 128.196ms = 5.381ms。

方向2: 使用Pin Memory

问题分析。
在方向1的基础上，继续寻找可优化的部分。从方向1的Timeline中可以看到，数据由Host端传入GPU端时（耗时约为13ms左右），GPU处于空闲状态，未做任何计算操作，那么缩短数据传输时间将有助于减少Batch的处理时间。

方案设计。

在传输数据时，尝试使用Pin Memory，在生产随机数据时，使用Pin Memory保存数据（由函数data_generation_with_pin_memory完成），main函数需做一定修改，其他代码基本不变，代码保存在3_use-pin-memory.py中。

展开查看完整代码

import os  # 导入用于操作系统接口的库
import nvtx  # 导入用于GPU性能分析的英伟达 Tools Extension库
import ctypes  # 导入 ctypes 模块，可以调用 C 语言库函数
import numpy as np  # 导入用于数值计算的NumPy库
import tensorrt as trt  # 导入TensorRT库
from cuda import cudart  # 导入CUDA Runtime API库

np.random.seed(10088)  # 设置NumPy的随机数种子以便重现结果


# 定义data_generation_with_pin_memory函数，用于生成随机数据并在主机端使用页锁定（pinned）内存，以加速主机到设备的数据传输
def data_generation_with_pin_memory(shape, batches):
    data = []  # 创建数据列表
    pbuffers = []  # 创建页锁定缓冲区列表
    for i in range(batches):  # 对于每个批次
        d = np.random.randn(*shape).astype(np.float32)  # 生成一批随机数据
        nElement = d.size  # 获取数组中元素个数
        nByteSize = d.nbytes  # 获取数据所占字节大小
        _, pBuffer = cudart.cudaHostAlloc(nByteSize, cudart.cudaHostAllocDefault)  # 使用CUDA API申请页锁定内存
        pBufferCtype = ctypes.cast(pBuffer, ctypes.POINTER(ctypes.c_float * nElement))  # 创建 ctypes 类型，方便与 CUDA API 交互
        nd = np.ndarray(shape=d.shape, dtype=d.dtype, buffer=pBufferCtype.contents)  # 创建 NumPy 数组用作页锁定内存的视图
        nd[:] = d  # 将生成的数据拷贝到页锁定内存中
        data.append(nd)  # 将页锁定内存数组加入数据列表
        pbuffers.append(pBuffer)  # 将页锁定内存指针加入列表
    return data, pbuffers  # 返回数据列表和页锁定缓冲区指针列表


# 定义main函数，它是程序的入口点
def main():
    data, pBuffers = data_generation_with_pin_memory([128, 3, 224, 224], 10)  # 使用页锁定内存生成10批次数据
    logger = trt.Logger(trt.Logger.ERROR)  # 创建TensorRT日志对象
    engine = load_engine(logger, "resnet18.plan")  # 加载TensorRT推理引擎
    infer(engine, data)  # 执行推理
    for p in pBuffers:  # 遍历并释放所有页锁定内存
        cudart.cudaFreeHost(p)


if __name__ == "__main__":
    main()

准备完成后，运行如下Shell代码。

nsys profile -w true \
	-t cuda,nvtx,osrt,cudnn,cublas \
	--cuda-memory-usage=true \
	--cudabacktrace=all \
	--cuda-graph-trace=node \
	--gpu-metrics-device=all \
	-f true \
	-o reports/3_use-pin-memory \
	python 3_use-pin-memory.py

运行完成后，在./reports目录下，生成一个名称为3_use-pin-memory.nsys-rep文件，将其导入Nsight Systems中，Timeline如下。从上图中可以看到：
- 最后三个Batch总共花费时间约为108.348ms，时间缩短128.196ms - 108.348ms = 19.848ms。
- Batch传输时间由13.429ms缩短为6.912ms。

方向3: 使用FP16（或INT8）精度

问题分析。
在优化方向2的Timeline中，计算每个Batch时间约为27.230ms，内存消耗4.7GB。
方案设计。
如果可以在编译模型时开启FP16精度（或INT8精度）等方式，则可缩短Batch计算时间。开启FP16只需在BuilderConfig中添加如下一行。
```
config.set_flag(trt.BuilderFlag.FP16)
```

在0_build.py脚本中，已指定一个选项（--fp16）用于编译时开启FP16模式，同时复制一份3_use-pin-memory.py并命名为4_use-fp16.py，执行如下Shell脚本。

python 0_build.py --fp16  # 开启FP16模式

nsys profile -w true \
	-t cuda,nvtx,osrt,cudnn,cublas \
	--cuda-memory-usage=true \
	--cudabacktrace=all \
	--cuda-graph-trace=node \
	--gpu-metrics-device=all \
	-f true \
	-o reports/4_use-fp16 \
	python 4_use-fp16.py

运行完成后，在./reports目录下，生成一个名称为4_use-fp16.nsys-rep文件，导入Nsight Systems中，Timeline如下。从图中可以看到：
- 三个Batch总的消耗时间由108.348ms 缩短为49.309ms。
- Batch计算时间由27.230ms缩短为7.957ms。
- GPU内存使用量由4.7GB减少为2.39 GB。
  重要
  生产实践中，还需要有一个校准过程，以保证模型量化后的结果正确性，具体请参考TensorRT官方文档。

方向4: 使用重叠数据传输和数据计算

问题分析。
当我们进行量化操作后，数据传输时间相比于数据计算时间已变得不可忽略。此时，单纯的缩短数据传输时间已经不可行了。

方案设计。

要完成数据传输和数据计算重叠的目标，需要借助CUDA Stream。

在代码中修改添加：

创建3个cuda stream，一个用于数据从Host到Device传输操作，一个用于数据计算和结果返回操作。
创建3个cuda event，用于cuda stream之间以及GPU与Host的同步操作。
预先传输第一个Batch的数据，然后在第一个Batch计算时，同时传输第二个Batch的数据，以此类推，当计算第二个Batch时，同时传输第三个Batch数据。

下面是完整代码，保存在5_multi-streams.py。

展开查看完整代码

import os  # 导入os模块，提供了操作系统相关的功能
import nvtx  # 导入英伟达 Tools Extension库，用于GPU性能分析
import time  # 导入time模块，提供时间相关的功能
import ctypes  # 导入ctypes模块，用于与C语言库交互
import numpy as np  # 导入NumPy库，用于进行高性能的数学运算
import tensorrt as trt  # 导入TensorRT库，用于进行深度学习推理优化
from cuda import cudart  # 从cuda模块导入cudart子模块，提供CUDA Runtime API的Python接口

np.random.seed(10088)  # 设置NumPy的随机数种子以保证随机操作的可重现性


# 定义data_generation_with_pin_memory函数，用于生成输入数据并使用页锁定内存来优化数据传输
def data_generation_with_pin_memory(shape, batches):
    data = []  # 创建空列表用于存储数据
    pbuffers = []  # 创建空列表用于存储页锁定内存的指针
    for i in range(batches):  # 循环生成指定数量的批次
        d = np.random.randn(*shape).astype(np.float32)  # 生成随机数数组
        nElement = d.size  # 计算数组中元素的总数
        nByteSize = d.nbytes  # 计算数组的字节大小
        _, pBuffer = cudart.cudaHostAlloc(nByteSize, cudart.cudaHostAllocDefault)  # 使用CUDA Runtime API分配页锁定内存
        pBufferCtype = ctypes.cast(pBuffer, ctypes.POINTER(ctypes.c_float * nElement))  # 将页锁定内存的指针转换为C类型
        # 创建NumPy数组，将页锁定内存作为其缓冲区
        nd = np.ndarray(shape=d.shape, dtype=d.dtype, buffer=pBufferCtype.contents)
        nd[:] = d  # 将生成的数据复制到页锁定内存中
        data.append(nd)  # 将页锁定内存中的数组添加到数据列表中
        pbuffers.append(pBuffer)  # 将页锁定内存的指针添加到列表中
    return data, pbuffers  # 返回生成的数据和页锁定内存指针列表


# 定义load_engine函数，用于加载TensorRT推理引擎
def load_engine(logger, plan_file):
    with open(plan_file, "rb") as plan:  # 以二进制读取方式打开序列化的推理引擎文件
        # 使用TensorRT的Runtime接口和提供的logger来反序列化推理引擎
        engine = trt.Runtime(logger).deserialize_cuda_engine(plan.read())
    return engine  # 返回反序列化后的推理引擎


# 定义get_io_tensors函数，用于获取推理引擎的输入输出张量信息
def get_io_tensors(engine):
    num_io_tensors = engine.num_io_tensors  # 获取引擎的输入输出张量总数
    io_tensor_names = [engine.get_tensor_name(i) for i in range(num_io_tensors)]  # 获取所有张量的名称
    # 计算输入张量数量
    num_input_io_tensors = [engine.get_tensor_mode(io_tensor_names[i]) for i in range(num_io_tensors)].count(trt.TensorIOMode.INPUT)
    return num_io_tensors, io_tensor_names, num_input_io_tensors  # 返回张量数量、名称和输入张量数量


# 定义infer函数，用于执行TensorRT推理
def infer(engine, data):
    context = engine.create_execution_context()  # 创建TensorRT的执行上下文
    # 创建两个CUDA流，一个用于数据拷贝，一个用于计算
    _, computeStream = cudart.cudaStreamCreate()
    _, copyStream = cudart.cudaStreamCreate()
    # 创建三个CUDA事件，用于同步操作
    _, inputConsumedEvent = cudart.cudaEventCreate()
    _, inputCopyedEvent = cudart.cudaEventCreate()
    _, outputReadyEvent = cudart.cudaEventCreate()
    num_io_tensors, io_tensor_names, num_input_io_tensors = get_io_tensors(engine)  # 获取IO张量信息
    context.set_input_shape(io_tensor_names[0], data[0].shape)  # 设置输入张量形状
    bufferH, bufferD = [], []  # 初始化主机和设备缓冲区列表
    bufferH.append(data[0])  # 将第一批数据添加到主机缓冲区列表
    for i in range(num_input_io_tensors, num_io_tensors):  # 为输出张量在主机端分配空间
        bufferH.append(np.empty(context.get_tensor_shape(io_tensor_names[i]), dtype=trt.nptype(engine.get_tensor_dtype(io_tensor_names[i]))))
    for i in range(num_io_tensors):  # 在设备端为每个张量分配空间
        bufferD.append(cudart.cudaMalloc(bufferH[i].nbytes)[1])
    for i in range(num_io_tensors):  # 设置执行上下文中的张量地址
        context.set_tensor_address(io_tensor_names[i], int(bufferD[i]))
    tet = None  # 初始化用于NVTX范围的变量
    # 预先传输第一个批次的输入数据
    for i in range(num_input_io_tensors):
        cudart.cudaMemcpyAsync(bufferD[i], bufferH[i].ctypes.data, bufferH[i].nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, copyStream)
    cudart.cudaEventRecord(inputCopyedEvent, copyStream)  # 标记输入数据拷贝完成的事件
    for i in range(len(data)):  # 对每批数据执行推理
        if i == 7:  # 如果是第7批数据
            tet = nvtx.start_range(message="Total Elapsed Time(3 batchs)", color="orange")  # 使用NVTX标记性能测试的开始
        d = None  # 初始化变量，用于指向后续批次的数据
        if i + 1 < len(data):  # 如果有后续批次的数据
            d = data[i + 1]
        nvtx.push_range(message="infer", color="purple")  # 使用NVTX标记单次推理的开始
        # 调用infer_once函数进行单次推理
        infer_once(engine, context, bufferH, bufferD, d, copyStream, computeStream, inputConsumedEvent, inputCopyedEvent, outputReadyEvent)
        nvtx.pop_range()  # 使用NVTX标记单次推理的结束
    nvtx.end_range(tet)  # 使用NVTX标记性能测试的结束
    for b in bufferD:  # 释放设备端的内存
        cudart.cudaFree(b)
    # 销毁CUDA流和事件
    cudart.cudaStreamDestroy(copyStream)
    cudart.cudaStreamDestroy(computeStream)
    cudart.cudaEventDestroy(inputConsumedEvent)
    cudart.cudaEventDestroy(inputCopyedEvent)
    cudart.cudaEventDestroy(outputReadyEvent)


# 定义infer_once函数，用于执行单次推理并处理输入输出数据
def infer_once(engine, context, bufferH, bufferD, data, copyStream, computeStream, inputConsumedEvent, inputCopyedEvent, outputReadyEvent):
    num_io_tensors, io_tensor_names, num_input_io_tensors = get_io_tensors(engine)  # 获取IO张量信息
    cudart.cudaEventSynchronize(inputCopyedEvent)  # 等待上一个输入数据拷贝完成
    if data is not None:  # 如果提供了新的数据
        bufferH[0] = data  # 更新主机缓冲区的数据
    # 在computeStream中等待inputCopyedEvent，保证上一个Batch输入数据已准备好
    cudart.cudaStreamWaitEvent(computeStream, inputCopyedEvent, cudart.cudaEventWaitDefault)
    context.execute_async_v3(computeStream)  # 在computeStream中异步执行推理计算
    # 标记输入数据已被消费的事件
    context.set_input_consumed_event(inputConsumedEvent)
    # 在copyStream中等待inputConsumedEvent，保证输入数据已被消费
    cudart.cudaStreamWaitEvent(copyStream, inputConsumedEvent, cudart.cudaEventWaitDefault)
    if bufferH[0] is not None:  # 如果有新的输入数据
        # 异步拷贝新的输入数据到设备端
        for i in range(num_input_io_tensors):
            cudart.cudaMemcpyAsync(bufferD[i], bufferH[i].ctypes.data, bufferH[i].nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, copyStream)
        cudart.cudaEventRecord(inputCopyedEvent, copyStream)  # 标记新的输入数据拷贝完成事件
    # 在computeStream中异步拷贝推理结果到主机端
    for i in range(num_input_io_tensors, num_io_tensors):
        cudart.cudaMemcpyAsync(bufferH[i].ctypes.data, bufferD[i], bufferH[i].nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, computeStream)
    cudart.cudaEventRecord(outputReadyEvent, computeStream)  # 标记推理结果拷贝完成事件
    cudart.cudaEventSynchronize(outputReadyEvent)  # 等待推理结果拷贝完成
    nvtx.push_range(message="Print Result")  # 使用NVTX标记打印结果的开始
    print_result(io_tensor_names, num_input_io_tensors, num_io_tensors, bufferH)  # 调用print_result函数打印推理结果
    nvtx.pop_range()  # 使用NVTX标记打印结果的结束
    print("Succeeded running model in TensorRT!")  # 打印推理成功的信息


# 定义print_result函数，用于输出推理结果
def print_result(io_tensor_names, num_input_io_tensors, num_io_tensors, bufferH):
    for i in range(num_input_io_tensors, num_io_tensors):  # 遍历所有输出张量
        softmax_scores = softmax(bufferH[i])  # 对输出张量应用softmax函数
        predicted_classes = np.argmax(softmax_scores, axis=1)  # 获取预测分类结果
        max_probs_np = np.max(softmax_scores, axis=1)  # 获取最高概率值
        print("Output Tensor Name: ", io_tensor_names[i])  # 打印输出张量名称
        print("Maximum probability for each image in the batch:\n", max_probs_np)  # 打印每张图片的最高概率
        print("Index of predicted class for each image in the batch:\n", predicted_classes)  # 打印预测分类索引


# 定义main函数，它是程序的入口点
def main():
    data, pBuffers = data_generation_with_pin_memory([128, 3, 224, 224], 10)  # 使用页锁定内存生成10批数据
    logger = trt.Logger(trt.Logger.ERROR)  # 创建TensorRT日志对象
    engine = load_engine(logger, "resnet18.plan")  # 加载TensorRT推理引擎
    infer(engine, data)  # 使用推理引擎执行推理
    for p in pBuffers:  # 遍历页锁定内存指针列表
        cudart.cudaFreeHost(p)  # 释放每个页锁定内存块


if __name__ == "__main__":
    main()  # 如果脚本作为主程序运行，则执行main函数

准备完成后，执行如下Shell代码。

python 0_build.py --fp16

nsys profile -w true \
	-t cuda,nvtx,osrt,cudnn,cublas \
	--cuda-memory-usage=true \
	--cudabacktrace=all \
	--cuda-graph-trace=node \
	--gpu-metrics-device=all \
	-f true \
	-o reports/5_multi-streams \
	python 5_multi-streams.py

结果在的Timeline中展示如下。
从图中可以获得如下信息：
- 3个Batch总耗时由49.309ms缩短为29.650ms。
- 在优化方向3的Timeline中，数据传输（Host端到GPU端，图中"Memcpy HtoD"）与数据计算（图中蓝色部分）是串行执行的，而在上图中数据传输和数据计算是并行执行的（图中标号3）。
- Batch与Batch之间还存在一定的空隙，从图中可以看出，这个是由print_result函数引起的。

方向5: 使用多线程处理输出结果

问题分析。
在方向4的Timeline中，Batch与Batch之间还存在一定的间隙，这些间隙是由print_result函数引起的，它在打印输出结果时，GPU是处于空闲状态的。

方案设计。

可以使用Python多线程另起一个单独的线程处理输出结果，主线程继续执行。

修改infer函数，创建一个线程执行print_result函数，完整代码如下，保存在6_multi-threads.py中。

展开查看完整代码

import nvtx  # 导入用于GPU性能分析的英伟达 Tools Extension库
import queue  # 导入Python标准库的队列模块
import threading  # 导入Python标准库的线程模块
import ctypes  # 导入ctypes模块，用于调用C语言库函数
import numpy as np  # 导入NumPy库，用于数值计算
import tensorrt as trt  # 导入TensorRT库
from cuda import cudart  # 从cuda模块导入cudart，提供CUDA Runtime API

np.random.seed(10088)  # 设置NumPy的随机数种子


# 定义softmax函数，用于执行神经网络的归一化
def softmax(x, axis=1):
    e_x = np.exp(x - np.max(x, axis=axis, keepdims=True))  # 减去最大值以提高数值稳定性
    return e_x / np.sum(e_x, axis=axis, keepdims=True)  # 利用广播机制归一化e_x

# 定义data_generation_with_pin_memory函数，用于生成随机数据并使用页锁定内存
def data_generation_with_pin_memory(shape, batches):
    data = []  # 创建数据列表
    pbuffers = []  # 创建页锁定内存缓冲区列表
    for i in range(batches):  # 对于每个批次
        d = np.random.randn(*shape).astype(np.float32)  # 生成一批随机数据
        nElement = d.size  # 获取数据中的元素总数
        nByteSize = d.nbytes  # 获取数据的总字节大小
        _, pBuffer = cudart.cudaHostAlloc(nByteSize, cudart.cudaHostAllocDefault)  # 调用CUDA API申请页锁定内存
        pBufferCtype = ctypes.cast(pBuffer, ctypes.POINTER(ctypes.c_float * nElement))  # 将页锁定内存的指针转换为ctypes类型
        nd = np.ndarray(shape=d.shape, dtype=d.dtype, buffer=pBufferCtype.contents)  # 使用numpy ndarray封装内存
        nd[:] = d  # 将数据复制到页锁定内存中
        data.append(nd)  # 将封装页锁定内存的numpy数组添加到数据列表中
        pbuffers.append(pBuffer)  # 将页锁定内存指针添加到pbuffers列表中
    return data, pbuffers  # 返回数据列表和页锁定内存指针列表


# 定义load_engine函数，用于加载TensorRT推理引擎
def load_engine(logger, plan_file):
    with open(plan_file, "rb") as plan:  # 以二进制读取模式打开推理引擎文件
        engine = trt.Runtime(logger).deserialize_cuda_engine(plan.read())  # 反序列化推理引擎
    return engine  # 返回反序列化后的推理引擎

# 定义get_io_tensors函数，获取推理引擎的输入输出张量信息
def get_io_tensors(engine):
    num_io_tensors = engine.num_io_tensors  # 获取引擎的输入输出张量总数
    io_tensor_names = [engine.get_tensor_name(i) for i in range(num_io_tensors)]  # 获取所有张量的名称
    num_input_io_tensors = [engine.get_tensor_mode(io_tensor_names[i]) for i in range(num_io_tensors)].count(trt.TensorIOMode.INPUT)  # 计算输入张量的数量
    return num_io_tensors, io_tensor_names, num_input_io_tensors  # 返回张量数量、名称和输入张量数量

# 定义infer函数，用于执行推理并创建线程以打印结果
def infer(engine, data):
    context = engine.create_execution_context()  # 创建TensorRT的执行上下文
    _, computeStream = cudart.cudaStreamCreate()  # 创建CUDA计算流
    _, copyStream = cudart.cudaStreamCreate()  # 创建CUDA拷贝流
    # 创建CUDA事件
    _, inputConsumedEvent = cudart.cudaEventCreate()
    _, inputCopyedEvent = cudart.cudaEventCreate()
    _, outputReadyEvent = cudart.cudaEventCreate()
    num_io_tensors, io_tensor_names, num_input_io_tensors = get_io_tensors(engine)  # 获取IO张量信息
    context.set_input_shape(io_tensor_names[0], data[0].shape)  # 设置输入张量的形状
    bufferH, bufferD = [], []  # 创建主机（Host）和设备（Device）缓冲区列表
    bufferH.append(data[0])  # 将第一批数据添加到主机缓冲区列表
    for i in range(num_input_io_tensors, num_io_tensors):  # 为输出张量在主机端分配空间
        bufferH.append(np.empty(context.get_tensor_shape(io_tensor_names[i]), dtype=trt.nptype(engine.get_tensor_dtype(io_tensor_names[i]))))
    for i in range(num_io_tensors):  # 在设备端为每个张量分配空间
        bufferD.append(cudart.cudaMalloc(bufferH[i].nbytes)[1])
    for i in range(num_io_tensors):  # 设置执行上下文中的张量地址
        context.set_tensor_address(io_tensor_names[i], int(bufferD[i]))
    q = queue.Queue(maxsize=1)  # 创建一个队列，用于线程间通信
    p = threading.Thread(target=print_result_in_thread, args=(q, io_tensor_names, num_input_io_tensors, num_io_tensors, bufferH))  # 创建打印结果的线程
    p.start()  # 启动线程
    tet = None  # 初始化NVTX范围变量
    # 预先传输第一个批次的输入数据
    for i in range(num_input_io_tensors):
        cudart.cudaMemcpyAsync(bufferD[i], bufferH[i].ctypes.data, bufferH[i].nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, copyStream)
    cudart.cudaEventRecord(inputCopyedEvent, copyStream)  # 记录输入数据拷贝完成事件
    for i in range(len(data)):  # 对每批数据执行推理
        if i == 7:  # 对于第7批数据进行特殊处理
            tet = nvtx.start_range(message="Total Elapsed Time(3 batchs)", color="orange")  # 使用NVTX开始记录时间范围
        d = None  # 初始化变量用于存储下一批数据
        if i + 1 < len(data):  # 如果存在下一批数据
            d = data[i + 1]
        nvtx.push_range(message="infer", color="purple")  # 使用NVTX标记推理过程
        infer_once(engine, context, bufferH, bufferD, d, copyStream, computeStream, inputConsumedEvent, inputCopyedEvent, outputReadyEvent, q)  # 调用infer_once进行单次推理
        nvtx.pop_range()  # 使用NVTX结束标记推理过程
    nvtx.end_range(tet)  # 使用NVTX结束记录时间范围
    q.put(None)  # 向队列放入None，标志结束
    for b in bufferD:  # 释放设备端内存
        cudart.cudaFree(b)
    # 销毁CUDA流和事件
    cudart.cudaStreamDestroy(copyStream)
    cudart.cudaStreamDestroy(computeStream)
    cudart.cudaEventDestroy(inputConsumedEvent)
    cudart.cudaEventDestroy(inputCopyedEvent)
    cudart.cudaEventDestroy(outputReadyEvent)

# 定义infer_once函数，用于执行单次推理操作
def infer_once(engine, context, bufferH, bufferD, data, copyStream, computeStream, inputConsumedEvent, inputCopyedEvent, outputReadyEvent, q):
    num_io_tensors, io_tensor_names, num_input_io_tensors = get_io_tensors(engine)  # 获取IO张量信息
    cudart.cudaEventSynchronize(inputCopyedEvent)  # 同步拷贝事件，确保输入数据已拷贝完成
    if data is not None:  # 如果提供了新数据
        bufferH[0] = data  # 更新主机端输入缓冲区的数据
    cudart.cudaStreamWaitEvent(computeStream, inputCopyedEvent, cudart.cudaEventWaitDefault)  # 让计算流等待输入数据拷贝完成
    context.execute_async_v3(computeStream)  # 在计算流中异步执行推理
    context.set_input_consumed_event(inputConsumedEvent)  # 设置输入数据消费完毕的事件
    cudart.cudaStreamWaitEvent(copyStream, inputConsumedEvent, cudart.cudaEventWaitDefault)  # 让拷贝流等待输入数据被消费
    if data is not None:  # 如果还有新数据
        bufferH[0] = data  # 更新主机端输入缓冲区的数据
        for i in range(num_input_io_tensors):  # 拷贝新的输入数据到设备端
            cudart.cudaMemcpyAsync(bufferD[i], bufferH[i].ctypes.data, bufferH[i].nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, copyStream)
        cudart.cudaEventRecord(inputCopyedEvent, copyStream)  # 记录输入数据拷贝完成的事件
    for i in range(num_input_io_tensors, num_io_tensors):  # 拷贝输出数据到主机端
        cudart.cudaMemcpyAsync(bufferH[i].ctypes.data, bufferD[i], bufferH[i].nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, computeStream)
    cudart.cudaEventRecord(outputReadyEvent, computeStream)  # 记录输出数据准备完毕的事件
    cudart.cudaEventSynchronize(outputReadyEvent)  # 同步输出数据准备完毕的事件，确保拷贝操作完成
    q.put("OutputReady")  # 向队列发送输出数据准备完毕的信号

# 定义print_result_in_thread函数，它将在一个单独的线程中运行，从队列中读取信号并打印推理结果
def print_result_in_thread(queue, io_tensor_names, num_input_io_tensors, num_io_tensors, bufferH):
    while True:  # 无限循环
        get_result = queue.get()  # 从队列读取项目，如果队列为空则阻塞
        if get_result is None:  # 如果收到了None信号，表示结束
            break
        nvtx.push_range(message="Print Result")  # 开始NVTX范围，用于性能分析
        print_result(io_tensor_names, num_input_io_tensors, num_io_tensors, bufferH)  # 调用print_result打印推理结果
        nvtx.pop_range()  # 结束NVTX范围

# 定义print_result函数，用于打印推理出的softmax结果
def print_result(io_tensor_names, num_input_io_tensors, num_io_tensors, bufferH):
    for i in range(num_input_io_tensors, num_io_tensors):  # 遍历所有输出张量
        softmax_scores = softmax(bufferH[i])  # 应用softmax函数计算概率
        predicted_classes = np.argmax(softmax_scores, axis=1)  # 使用np.argmax找到概率最高的类别
        max_probs_np = np.max(softmax_scores, axis=1)  # 找到最高概率值
        print("Output Tensor Name: ", io_tensor_names[i])  # 打印输出张量的名称
        print("Maximum probability for each image in the batch:\n", max_probs_np)  # 打印每张图片的最高概率
        print("Index of predicted class for each image in the batch:\n", predicted_classes)  # 打印预测的类别索引


# 定义main函数，作为程序的入口点
def main():
    data, pBuffers = data_generation_with_pin_memory([128, 3, 224, 224], 10)  # 生成输入数据并使用页锁定内存
    logger = trt.Logger(trt.Logger.ERROR)  # 创建TensorRT日志对象
    engine = load_engine(logger, "resnet18.plan")  # 加载TensorRT推理引擎
    infer(engine, data)  # 执行推理过程
    for p in pBuffers:  # 遍历页锁定内存指针列表
        cudart.cudaFreeHost(p)  # 释放页锁定内存

if __name__ == "__main__":
    main()  # 如果脚本作为主程序运行，则执行main函数

准备完成后，执行如下Shell代码。

python 0_build.py --fp16

nsys profile -w true \
	-t cuda,nvtx,osrt,cudnn,cublas \
	--cuda-memory-usage=true \
	--cudabacktrace=all \
	--cuda-graph-trace=node \
	--gpu-metrics-device=all \
	-f true \
	-o reports/6_multi-threads \
	python 6_multi-threads.py

在的Timeline显示结果如下，可以看到print_result函数与GPU计算与数据传输重叠，并不是串行执行的。同时，最后三个Batch的耗时缩短29.650ms - 26.787ms = 2.86ms。

方向6: 使用CUDA Graph

问题分析。
在方向5的Timeline中，每次Batch计算都会出现25次Kernel Launch，每次Kernel Launch都会消耗一定的时间。

方案设计。

CUDA Graphs 是 CUDA 10.0 中引入的一个特性，它允许开发者捕获一系列CUDA操作（如内存传输和核函数执行），并将它们组织成一个被称为“graph”的有向无环图。这个图可以被看作是在不同CUDA流上执行的一系列操作的“快照”，一旦捕获，就可以多次执行，而无需CPU介入，从而减少了CPU与GPU之间的交互，提高了整体的执行效率。可以利用CUDA Graph来捕获TensorRT引擎执行推理所需要的一系列CUDA操作，包括内存拷贝、核函数执行等。这样可以减少CPU参与的推理调度开销，尤其是在执行重复推理任务时，可以显著提高推理性能。

重要

请注意CUDA Graphs的使用可能具有一定的复杂性，通常需要对CUDA编程有一定的了解，以及对如何在TensorRT中配置和执行推理过程有深入的理解。此外，CUDA Graphs的有效性也取决于应用场景，不是所有的应用都能从中获益，因为有些CUDA 操作不一定支持CUDA Graph。本文涉及到的模型就是一个例子，无法使用CUDA Graph，因为模型中存在不支持CUDA Graph的操作。

在TensorRT中使用CUDA Graph的示例如下，修改infer_once函数。

展开查看完整代码

# 定义一个名为infer_once的函数，它将执行单次推理，使用TensorRT引擎和CUDA流、事件
def infer_once(engine,context,bufferH,bufferD,data,copyStream,computeStream,inputConsumedEvent,inputCopyedEvent,outputReadyEvent,q):
    # 调用get_io_tensors函数获取TensorRT引擎的输入输出张量的数量、名称以及输入张量的数量。
	num_io_tensors,io_tensor_names,num_input_io_tensors = get_io_tensors(engine)
    # 同步inputCopyedEvent事件，确保之前的数据复制操作已经完成。
	cudart.cudaEventSynchronize(inputCopyedEvent)
    # 将新的数据赋值给主机缓冲区bufferH的第一个元素。
	bufferH[0] = data
	# 在computeStream中等待数据传输完成
	cudart.cudaStreamWaitEvent(computeStream, inputCopyedEvent, cudart.cudaEventWaitDefault)
    # 开始在计算流computeStream上捕获操作，以便创建CUDA Graph。此处指定的捕获模式为全局模式。
	cudart.cudaStreamBeginCapture(computeStream, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeGlobal)
	# 在计算流上异步执行推理计算。这是CUDA Graph捕获过程的一部分。
    context.execute_async_v3(computeStream)
    # 结束捕获计算流，并创建一个CUDA Graph对象。打印状态信息，以确认操作是否成功。 
	status, graph = cudart.cudaStreamEndCapture(computeStream)
    # 打印状态信息
	print(status)
    # 实例化CUDA Graph，准备执行，并打印状态信息。
	status, graphExec = cudart.cudaGraphInstantiate(graph, 0)
    # 打印状态
	print(status)
    # 在计算流上启动CUDA Graph执行。这将执行前面捕获的所有CUDA操作。
	cudart.cudaGraphLaunch(graphExec, computeStream)
	# 设置输入数据被消费的标记
	context.set_input_consumed_event(inputConsumedEvent)

	# 计算完成后，数据传回
	for i in range(num_input_io_tensors, num_io_tensors):
		cudart.cudaMemcpyAsync(bufferH[i].ctypes.data, bufferD[i], bufferH[i].nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost,computeStream)
	cudart.cudaEventRecord(outputReadyEvent, computeStream)

	# 在copyStream中等待输入数据被消费
	cudart.cudaStreamWaitEvent(copyStream, inputConsumedEvent, cudart.cudaEventWaitDefault)
	# 传输后一个数据
	if data is not None:
		bufferH[0] = data
		for i in range(num_input_io_tensors):
			cudart.cudaMemcpyAsync(bufferD[i], bufferH[i].ctypes.data, bufferH[i].nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost,copyStream)
		cudart.cudaEventRecord(inputCopyedEvent, copyStream)

	cudart.cudaEventSynchronize(outputReadyEvent)
	q.put("OutputReady")

执行代码后，报如下错误，估计是模型中存在不支持图捕获的CUDA操作。关于cuda graph例子，可以参考Cuda Graph。

方向7: 使用多CPU线程多Context

问题分析。
上述的优化方向演示的都是单个Context处理Batch数据，单个Context无法并行处理多个Batch，原因在于每个Context都存在推理时的中间结果缓存，而这些缓存同一时间只能为一个Batch提供服务。

方案设计。

那么，有没有可能Host端有多个CPU线程提交Batch处理请求，并且TensorRT能够并行处理这些Batch数据呢？答案是可以的，TensorRT支持创建多个Context，各个Context之间不互相影响。

下面的例子将创建两个Context（两个Context使用同一个Profile，需要在编译模型时开启共享Profile的选项），然后产生的10个Batch数据分成两组，每组5个Batch数据。每个Context处理一组Batch。

完整代码如下，保存在8_multi_cpu-threads.py。

展开查看完整代码

import nvtx  # 导入英伟达 Tools Extension（NVTX）库，用于性能分析
import queue  # 导入Python的queue模块，用于线程间通信
import threading  # 导入Python的threading模块，用于多线程编程
import ctypes  # 导入ctypes模块，用于在Python中调用C语言库
import numpy as np  # 导入NumPy库，一种用于数值计算的Python库
import tensorrt as trt  # 导入TensorRT库，用于深度学习推断优化
from cuda import cudart  # 从cuda模块导入cudart，提供CUDA Runtime API的Python接口


np.random.seed(10088)


def softmax(x, axis=1):
    e_x = np.exp(x - np.max(x, axis=axis, keepdims=True)) # 为了数值稳定性
    return e_x / np.sum(e_x, axis=axis, keepdims=True)


def data_generation_with_pin_memory(shape, batches):
    data1, data2 = [], []  # 创建两个列表存储数据，分别对应两个不同的数据集
    pbuffers = []  # 创建一个列表以存储页锁定内存的指针

    for i in range(batches):  # 为每个批次生成和存储数据
        d = np.random.randn(*shape).astype(np.float32)  # 生成随机数据
        nElement = d.size  # 获取数组中元素的总数
        nByteSize = d.nbytes  # 获取数组的字节大小
        _, pBuffer = cudart.cudaHostAlloc(nByteSize, cudart.cudaHostAllocDefault)  # 使用CUDA分配页锁定内存
        pBufferCtype = ctypes.cast(pBuffer, ctypes.POINTER(ctypes.c_float * nElement))  # 将页锁定内存转换为ctypes类型
        nd = np.ndarray(shape=d.shape, dtype=d.dtype, buffer=pBufferCtype.contents)  # 创建NumPy数组封装页锁定内存
        nd[:] = d  # 将生成的数据复制到页锁定内存中
        if i % 2 == 0:  # 交替地将数据添加到两个不同的数据列表中
            data1.append(nd)
        else:
            data2.append(nd)
        pbuffers.append(pBuffer)  # 将页锁定内存指针添加到列表中

    return data1, data2, pbuffers  # 返回两组数据列表和页锁定内存指针列表


def load_engine(logger,plan_file):
	with open(plan_file, "rb") as plan:
			engine = trt.Runtime(logger).deserialize_cuda_engine(plan.read())
	return engine


def get_io_tensors(engine):
	num_io_tensors = engine.num_io_tensors
	io_tensor_names = [engine.get_tensor_name(i) for i in range(num_io_tensors)]
	num_input_io_tensors = [engine.get_tensor_mode(io_tensor_names[i]) for i in range(num_io_tensors)].count(trt.TensorIOMode.INPUT)
	return num_io_tensors,io_tensor_names,num_input_io_tensors


def infer(engine,data):
	context = engine.create_execution_context()
	# 创建两个cuda stream
	_, computeStream = cudart.cudaStreamCreate()
	_, copyStream = cudart.cudaStreamCreate()

	# 创建三个event
	_, inputConsumedEvent = cudart.cudaEventCreate()
	_, inputCopyedEvent = cudart.cudaEventCreate()
	_, outputReadyEvent = cudart.cudaEventCreate()

	num_io_tensors,io_tensor_names,num_input_io_tensors = get_io_tensors(engine)
	context.set_input_shape(io_tensor_names[0], data[0].shape)
	bufferH,bufferD = [],[]
	bufferH.append(data[0])
	for i in range(num_input_io_tensors, num_io_tensors):
		bufferH.append(np.empty(context.get_tensor_shape(io_tensor_names[i]), dtype=trt.nptype(engine.get_tensor_dtype(io_tensor_names[i]))))

	for i in range(num_io_tensors):
		bufferD.append(cudart.cudaMalloc(bufferH[i].nbytes)[1])

	for i in range(num_io_tensors):
		context.set_tensor_address(io_tensor_names[i], int(bufferD[i]))

	q = queue.Queue(maxsize=1)
	p = threading.Thread(target=print_result_in_thread,args=(q,io_tensor_names,num_input_io_tensors,num_io_tensors,bufferH))
	p.start()

	tet = None

	# 首先传输第一个数据
	for i in range(num_input_io_tensors):
		cudart.cudaMemcpyAsync(bufferD[i], bufferH[i].ctypes.data, bufferH[i].nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost,copyStream)
	cudart.cudaEventRecord(inputCopyedEvent, copyStream)

	for i in range(len(data)):
		if i == 2:
			tet = nvtx.start_range(message="Total Elapsed Time(3 batchs)", color="orange")
		d = None
		if i+1 <= len(data) -1:
			d = data[i+1]
		nvtx.push_range(message="infer",color="purple")
		infer_once(engine,context,bufferH,bufferD,d,copyStream,computeStream,inputConsumedEvent,inputCopyedEvent,outputReadyEvent,q)
		nvtx.pop_range()
	nvtx.end_range(tet)
	q.put(None)
	for b in bufferD:
		cudart.cudaFree(b)
	cudart.cudaStreamDestroy(copyStream)
	cudart.cudaStreamDestroy(computeStream)
	cudart.cudaEventDestroy(inputConsumedEvent)
	cudart.cudaEventDestroy(inputCopyedEvent)
	cudart.cudaEventDestroy(outputReadyEvent)

def infer_once(engine,context,bufferH,bufferD,data,copyStream,computeStream,inputConsumedEvent,inputCopyedEvent,outputReadyEvent,q):
	num_io_tensors,io_tensor_names,num_input_io_tensors = get_io_tensors(engine)
	cudart.cudaEventSynchronize(inputCopyedEvent)
	bufferH[0] = data
	# 在computeStream中等待数据传输完成
	cudart.cudaStreamWaitEvent(computeStream, inputCopyedEvent, cudart.cudaEventWaitDefault)
	context.execute_async_v3(computeStream)
	# 设置输入数据被消费的标记
	context.set_input_consumed_event(inputConsumedEvent)
	# 在copyStream中等待输入数据被消费
	cudart.cudaStreamWaitEvent(copyStream, inputConsumedEvent, cudart.cudaEventWaitDefault)
	# 传输后一个数据
	if data is not None:
		bufferH[0] = data
		for i in range(num_input_io_tensors):
			cudart.cudaMemcpyAsync(bufferD[i], bufferH[i].ctypes.data, bufferH[i].nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost,copyStream)
		cudart.cudaEventRecord(inputCopyedEvent, copyStream)

	# 计算完成后，数据传回
	for i in range(num_input_io_tensors, num_io_tensors):
		cudart.cudaMemcpyAsync(bufferH[i].ctypes.data, bufferD[i], bufferH[i].nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost,computeStream)
	cudart.cudaEventRecord(outputReadyEvent, computeStream)

	cudart.cudaEventSynchronize(outputReadyEvent)
	q.put("OutputReady")


def print_result_in_thread(queue,io_tensor_names,num_input_io_tensors,num_io_tensors,bufferH):
	while True:
		get_result = queue.get()
		if get_result is None:
			break
		nvtx.push_range(message="Print Result")
		print_result(io_tensor_names,num_input_io_tensors,num_io_tensors,bufferH)
		nvtx.pop_range()


def print_result(io_tensor_names,num_input_io_tensors,num_io_tensors,bufferH):
	for i in range(num_input_io_tensors,num_io_tensors):
		softmax_scores = softmax(bufferH[i])
		# 使用np.argmax函数获取概率最高的类别
		predicted_classes = np.argmax(softmax_scores, axis=1)
		max_probs_np = np.max(softmax_scores, axis=1)
		print("Output Tensor Name: ",io_tensor_names[i])
		print("Maximum probability for each image in the batch:\n", max_probs_np)
		print("Index of predicted class for each image in the batch:\n", predicted_classes)


def main():
    data1, data2, pBuffers = data_generation_with_pin_memory([128, 3, 224, 224], 10)  # 使用页锁定内存生成两组数据
    logger = trt.Logger(trt.Logger.ERROR)  # 创建TensorRT日志对象
    engine = load_engine(logger, "resnet18.plan")  # 加载TensorRT推理引擎

    threads = []  # 创建线程列表
    p = threading.Thread(target=infer, args=(engine, data1))  # 创建一个新的线程来执行第一组数据的推理
    threads.append(p)  # 将线程添加到列表中
    p = threading.Thread(target=infer, args=(engine, data2))  # 创建一个新的线程来执行第二组数据的推理
    threads.append(p)  # 将线程添加到列表中

    for p in threads:  # 遍历线程列表
        p.start()  # 启动每个线程
    for p in threads:  # 遍历线程列表
        p.join()  # 等待每个线程完成

    for p in pBuffers:  # 遍历页锁定内存指针列表
        cudart.cudaFreeHost(p)  # 释放页锁定内存


if __name__ == "__main__":
	main()

在上述代码中，启动两个Python线程，每个线程处理5个Batch，每个线程将创建一个Context，执行各自的推理。准备完成后，执行如下Shell脚本。

python 0_build.py --fp16 --share-profile # 开启共享Profile模式

nsys profile -w true \
	-t cuda,nvtx,osrt,cudnn,cublas \
	--cuda-memory-usage=true \
	--cudabacktrace=all \
	--cuda-graph-trace=node \
	--gpu-metrics-device=all \
	-f true \
	-o reports/8_multi-context \
	python 8_multi-context.py

在Timeline中显示结果如下。
可以得到如下结论：
- GPU计算在两个Stream（Stream 19和Stream 16）中并行执行。
- 每个Stream中最后3个Batch的消耗时间约为50ms左右，最后三个Batch平均使用时间约为50 / 2 = 25ms，每个Batch计算时间约为13ms。与单个Batch相比计算时间有所增加，说明GPU资源是充分利用，但有可能造成计算任务过载的问题。
- GPU内存使用量是单个Context执行推理时的2倍左右。

方向8: 使用Dynamic Shape多重配置文件

问题分析。
使用动态形状（Dynamic Shape）来提高内核的灵活性，从而可以处理不同大小的输入数据。然而，在实际使用中，我们可能需要为不同的输入大小或其他条件编译多个版本的内核，这就是多Profile（多重配置文件）。多Profile可以通过在编译内核时添加多个编译选项来实现。例如，我们可以为不同的输入大小使用不同的线程块大小或其他优化参数。某些场景下，Dynamic Shape模式在min-opt-max跨度较大时，性能下降比较明显。

方案设计。

解决办法就是使用多个OptimizationProfile，对应多个min-opt-max，并且缩小min-opt-max跨度。

为了演示多Profile的使用过程，代码将基于Baseline的代码做简化修改。

首先修改0_build.py代码，需要在编译阶段构建两个Profile，并为每个Profile设置min-opt-max shape，保存在9_multi-profiles-build.py。

展开查看完整代码

import argparse
import os
import tensorrt as trt


def build(logger,ONNX_file,shapes,num_aux_stream,share_profile,fp16):
    errors = []  # 初始化一个空列表来收集过程中的错误信息
    builder = trt.Builder(logger)  # 创建一个TensorRT Builder实例
    # 创建支持显式批处理的网络定义
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    # 为每个输入形状创建一个优化配置文件
    profiles = [builder.create_optimization_profile() for _ in range(len(shapes))]
    config = builder.create_builder_config()  # 创建构建器配置对象
    if share_profile:  # 如果启用共享配置文件
        print("enable share profile")
        # 启用共享优化配置文件的预览特性
        config.set_preview_feature(trt.PreviewFeature.PROFILE_SHARING_0806, True)
    if num_aux_stream > 0:  # 如果有附加流
        print("set aux stream " + str(num_aux_stream))
        config.max_aux_streams = num_aux_stream  # 设置附加流的数量
    if fp16:  # 如果使用FP16精度
        config.set_flag(trt.BuilderFlag.FP16)  # 启用FP16精度模式
    parser = trt.ONNXParser(network, logger)  # 创建ONNX解析器
    if not os.path.exists(ONNX_file):  # 检查ONNX文件是否存在
        errors.append("Failed to find ONNX File!")  # 如果不存在，添加错误信息
        return None, errors  # 返回None和错误信息列表
    with open(ONNX_file, "rb") as model:  # 以二进制读模式打开ONNX文件
        if not parser.parse(model.read()):  # 读取并解析ONNX模型
            errors.append("failed to parse .ONNX file: ")  # 如果解析失败，添加错误信息
            for error in range(parser.num_errors):  # 遍历所有解析错误
                errors.append(parser.get_error(error))  # 添加每个解析错误信息
            return None, errors  # 返回None和错误信息列表
    inputTensor = network.get_input(0)  # 获取网络的第一个输入张量
    for i in range(len(shapes)):  # 遍历所有输入形状
        # 为每个输入张量设置优化配置文件的输入形状
        profiles[i].set_shape(inputTensor.name, *shapes[i])
        config.add_optimization_profile(profiles[i])  # 将优化配置文件添加到构建器配置中
    config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED  # 设置详细的性能分析级别
    # 构建并序列化网络为引擎字符串
    engine_string = builder.build_serialized_network(network, config)
    if engine_string == None:  # 检查是否成功构建引擎
        errors.append("Failed to build engine")  # 如果构建失败，添加错误信息
        return None, errors  # 返回None和错误信息列表
    return engine_string, errors  # 如果构建成功，返回引擎字符串和错误信息列表


def save_engine(engine_string,planFile):
	with open(planFile, "wb") as f:
		f.write(engine_string)
	return 0


def main():
    # 使用argparse库创建命令行参数解析器
    parser = argparse.ArgumentParser(description='ResNet18 TensorRT Builder')
    # 添加命令行参数 --aux-stream，用来指定辅助流的数量，默认值为0
    parser.add_argument('--aux-stream', type=int, default=0, metavar='N',
                        help='specify the aux stream (default: 0)')
    # 添加命令行参数 --share-profile，用来启用共享配置文件，为一个开关选项，默认关闭
    parser.add_argument('--share-profile', action='store_true', default=False,
                        help='enable share profile')
    # 添加命令行参数 --fp16，用来启用FP16模式，为一个开关选项，默认关闭
    parser.add_argument('--fp16', action='store_true', default=False,
                        help='enable fp16 mode')
    # 添加命令行参数 --output，用来指定输出的计划文件(plan file)，默认值为'resnet18.plan'
    parser.add_argument('--output', type=str, default='resnet18.plan', metavar='N',
                        help='specify the plan file')
    # 添加命令行参数 --ONNX-file，用来指定输入的ONNX文件，默认值为'resnet18.ONNX'
    parser.add_argument('--ONNX-file', type=str, default='resnet18.ONNX', metavar='N',
                        help='specify the ONNX file')
    # 解析命令行输入的参数
    args = parser.parse_args()
    
    # 创建TensorRT日志对象，设置日志级别为ERROR
    logger = trt.Logger(trt.Logger.ERROR)
    # 定义网络输入张量的不同形状，这里定义了两个区间的形状，用于优化配置文件
    shapes = [
        [[1,3,224,224],[16,3,224,224],[32,3,224,224]],
        [[64,3,224,224],[128,3,224,224],[192,3,224,224]],
    ]
    
    # 调用build函数构建TensorRT引擎，传入日志对象、ONNX文件、形状、额外参数
    engine_string, errors = build(logger, args.ONNX_file, shapes, args.aux_stream, args.share_profile, args.fp16)
    
    # 检查build函数是否返回错误信息
    if len(errors) != 0:
        print(errors)  # 如果有错误，打印错误信息
        return 1  # 返回错误码1，表示程序异常退出
    # 如果构建成功，调用save_engine函数保存序列化的引擎到指定的文件
    save_engine(engine_string, args.output)
    return 0  # 正常退出程序，返回码为0

if __name__ == "__main__":
	main()

然后准备执行推理的代码，保存在9_multi-profiles.py。在代码中准备了两个数据集（data0和data1），shape分别为[16,3,224,224]和[128,3,224,224]。如果不开启多Profile，那么每次都需要为数据分配新的GPU内存。

展开查看完整代码

import argparse
import os
import tensorrt as trt


def build(logger,ONNX_file,shapes,num_aux_stream,share_profile,fp16):
    errors = []  # 初始化一个空列表来收集过程中的错误信息
    builder = trt.Builder(logger)  # 创建一个TensorRT Builder实例
    # 创建支持显式批处理的网络定义
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    # 为每个输入形状创建一个优化配置文件
    profiles = [builder.create_optimization_profile() for _ in range(len(shapes))]
    config = builder.create_builder_config()  # 创建构建器配置对象
    if share_profile:  # 如果启用共享配置文件
        print("enable share profile")
        # 启用共享优化配置文件的预览特性
        config.set_preview_feature(trt.PreviewFeature.PROFILE_SHARING_0806, True)
    if num_aux_stream > 0:  # 如果有附加流
        print("set aux stream " + str(num_aux_stream))
        config.max_aux_streams = num_aux_stream  # 设置附加流的数量
    if fp16:  # 如果使用FP16精度
        config.set_flag(trt.BuilderFlag.FP16)  # 启用FP16精度模式
    parser = trt.ONNXParser(network, logger)  # 创建ONNX解析器
    if not os.path.exists(ONNX_file):  # 检查ONNX文件是否存在
        errors.append("Failed to find ONNX File!")  # 如果不存在，添加错误信息
        return None, errors  # 返回None和错误信息列表
    with open(ONNX_file, "rb") as model:  # 以二进制读模式打开ONNX文件
        if not parser.parse(model.read()):  # 读取并解析ONNX模型
            errors.append("failed to parse .ONNX file: ")  # 如果解析失败，添加错误信息
            for error in range(parser.num_errors):  # 遍历所有解析错误
                errors.append(parser.get_error(error))  # 添加每个解析错误信息
            return None, errors  # 返回None和错误信息列表
    inputTensor = network.get_input(0)  # 获取网络的第一个输入张量
    for i in range(len(shapes)):  # 遍历所有输入形状
        # 为每个输入张量设置优化配置文件的输入形状
        profiles[i].set_shape(inputTensor.name, *shapes[i])
        config.add_optimization_profile(profiles[i])  # 将优化配置文件添加到构建器配置中
    config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED  # 设置详细的性能分析级别
    # 构建并序列化网络为引擎字符串
    engine_string = builder.build_serialized_network(network, config)
    if engine_string == None:  # 检查是否成功构建引擎
        errors.append("Failed to build engine")  # 如果构建失败，添加错误信息
        return None, errors  # 返回None和错误信息列表
    return engine_string, errors  # 如果构建成功，返回引擎字符串和错误信息列表


def save_engine(engine_string,planFile):
	with open(planFile, "wb") as f:
		f.write(engine_string)
	return 0


def main():
    # 使用argparse库创建命令行参数解析器
    parser = argparse.ArgumentParser(description='ResNet18 TensorRT Builder')
    # 添加命令行参数 --aux-stream，用来指定辅助流的数量，默认值为0
    parser.add_argument('--aux-stream', type=int, default=0, metavar='N',
                        help='specify the aux stream (default: 0)')
    # 添加命令行参数 --share-profile，用来启用共享配置文件，为一个开关选项，默认关闭
    parser.add_argument('--share-profile', action='store_true', default=False,
                        help='enable share profile')
    # 添加命令行参数 --fp16，用来启用FP16模式，为一个开关选项，默认关闭
    parser.add_argument('--fp16', action='store_true', default=False,
                        help='enable fp16 mode')
    # 添加命令行参数 --output，用来指定输出的计划文件(plan file)，默认值为'resnet18.plan'
    parser.add_argument('--output', type=str, default='resnet18.plan', metavar='N',
                        help='specify the plan file')
    # 添加命令行参数 --ONNX-file，用来指定输入的ONNX文件，默认值为'resnet18.ONNX'
    parser.add_argument('--ONNX-file', type=str, default='resnet18.ONNX', metavar='N',
                        help='specify the ONNX file')
    # 解析命令行输入的参数
    args = parser.parse_args()
    
    # 创建TensorRT日志对象，设置日志级别为ERROR
    logger = trt.Logger(trt.Logger.ERROR)
    # 定义网络输入张量的不同形状，这里定义了两个区间的形状，用于优化配置文件
    shapes = [
        [[1,3,224,224],[16,3,224,224],[32,3,224,224]],
        [[64,3,224,224],[128,3,224,224],[192,3,224,224]],
    ]
    
    # 调用build函数构建TensorRT引擎，传入日志对象、ONNX文件、形状、额外参数
    engine_string, errors = build(logger, args.ONNX_file, shapes, args.aux_stream, args.share_profile, args.fp16)
    
    # 检查build函数是否返回错误信息
    if len(errors) != 0:
        print(errors)  # 如果有错误，打印错误信息
        return 1  # 返回错误码1，表示程序异常退出
    # 如果构建成功，调用save_engine函数保存序列化的引擎到指定的文件
    save_engine(engine_string, args.output)
    return 0  # 正常退出程序，返回码为0

if __name__ == "__main__":
	main()

准备完成后，执行如下Shell脚本。

python 9_multi-profiles-build.py --output resnet18-multi-profiles.plan

nsys profile -w true \
	-t cuda,nvtx,osrt,cudnn,cublas \
	--cuda-memory-usage=true \
	--cudabacktrace=all \
	--cuda-graph-trace=node \
	--gpu-metrics-device=all \
	-f true \
	-o reports/9_multi-profiles \
	python 9_multi-profiles.py

将结果导入Nsight Systems中，效果如下。可以看到，前面10个Batch的shape为[16,3,224,224]。
后面10个Batch的shape为[128,3,224,224]。
关于dynamic shape更多的用法，可以参考TensorRT Cookbook。

总结

经过上述一系列的优化技巧取得了性能的提升，我们成功地将3个Batch的处理时间由133.577ms缩短为25ms左右。更多关于TensorRT的高级功能和优化技巧，可参考TensorRT Cookbook。