TensorRT 使用指南（1）：基本流程
安装 TensorRT 的注意事项
TensorRT 的安装方式在官方安装导引 里已经有详细的说明了，这里提下需要注意的地方，首先是尽量保证 TensorRT 的安装方式 和 CUDA 的安装方式相同，否则可能会出现找不到 CUDA 的情况，比如在 Ubuntu 中，如果 CUDA 不是通过 deb 包安装的，后面用 deb 安装 TensorRT 就会报错找不到 CUDA 依赖。
从 PyTorch 导出 ONNX 模型
ONNX 是一种通用的神经网络模型交换格式，TensorRT 有专门的 ONNX 解析器，可以解析 ONNX 模型并创建 TensorRT 自有的网络结构并做后续的优化工作。因此，为了让 TensorRT 优化我们使用 PyTorch 训练的模型，可以先将 PyTorch 模型导出为 ONNX 模型。值得注意的是，PyTorch 生态提供了 Torch-TensorRT 工具来和 TensorRT 交互，但这又是另外一种技术栈，目前我们先不考虑。
根据模型输入的不同， TensorRT 将输入分为两个类别，一个是固定张量形状的 static shape，和非固定张量形状的 dynamic shape。两种输入类别的导出方式有一些差异，首先 看一下 static shape，以 ResNet50 模型为例，其导出代码如下
import torch
from torchvision import models
model = models.resnet50(models.ResNet50_Weights.IMAGENET1K_V2)
model.eval()
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(model,
dummy_input,
"./weights/model.onnx",
do_constant_folding=True,
verbose=True,
input_names=["input"],
output_names=["output"])
这里的 dummy_input 就是输入样例，ONNX 将跟踪 PyTorch 模型的计算过程，生成计算图，然后导出为 ONNX 模型。
接下来我们以 Bert 模型为例，展示 dymamic shape 的导出方式
import os
import torch
from transformers import BertModel, BertTokenizer
model_path = "bert-base-chinese"
model = BertModel.from_pretrained(model_path)
tokenizer = BertTokenizer.from_pretrained(model_path)
onnx_export_path = os.path.join(model_path, "model.onnx")
max_seq_len = 512
model.eval()
dummpy_input = tokenizer("今天天气不错", return_tensors="pt", padding="max_length", truncation=True, max_length=max_seq_len)
with torch.no_grad():
symbolic_names = {0: 'batch_size', 1: 'max_seq_len'}
torch.onnx.export(model,
args=tuple(dummpy_input.values()),
f = onnx_export_path,
opset_version=13,
do_constant_folding=True,
input_names=["input_ids", "attention_mask", "token_type_ids"],
output_names=["last_hidden_state", "pooler_output"],
dynamic_axes={
'input_ids': symbolic_names,
'attention_mask': symbolic_names,
'token_type_ids': symbolic_names,
'last_hidden_state': symbolic_names,
'pooler_output': symbolic_names
})
这里的 dummpy_input 是一个字典类型，包含 input_ids，attention_mask 和 token_type_ids 三个值，它们被转换为 tuple 并作为 args 参数传入。需要注意 args 中元素的顺序必须和 input_names 参数指定的顺序保持一致，否则参数名和值就匹配错了。最后的 dynamic_axes 参数指定了哪些输入张量的形状是动态的，显然这里所有输入输出张量的第一个和第二个维度都是动态的，其中第一个维度是 batch，第二个维度是序列长度，而第三个维度是 embedding 维度，对于特定模型来说是固定的，所以这里的 symbolic_names 只包含了前两个维度。
ONNX 模型转换为 TensorRT Engine
ONNX 转换为 TensorRT Engine 至少要用到以下几个组件，logger，builder, parser, config 和 network。其中
- logger 负责日志记录，
- builder 负责多个对象的创建工作，
- parser 负责解析模型文件,
- config 负责保存配置信息，
- network 负责表示 TensorRT 的网络结构。
下面是一个最基础的 static_shape 例子
import tensorrt as trt
onnx_path = "./weights/model.onnx"
engine_path = "./weights/model.engine"
logger = trt.Logger(trt.Logger.VERBOSE)
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, logger)
config = builder.create_builder_config()
with open(onnx_path, "rb") as f:
parser.parse(f.read())
engineString = builder.build_serialized_network(network, config)
with open(engine_path, "wb") as f:
f.write(engineString)
流程相当清晰，首先创建 logger, builder, network, parser, config 等组件，然后读取模型的二进制数据，使用 parser 进行解析，由于 parser 包含 network 组件，所以它能把模型结构保存到 network 中，最后使用 builder 将 network 构建为序列化模型对象并存储。
如果是 dynamic shape 模型，则必须设置 optimization profile，也就是每个 input 的 shape
batch_size=1
max_seq_length=512
profile.set_shape("input_ids",
min = (batch_size, 1),
opt = (batch_size, max_seq_length//2),
max = (batch_size, max_seq_length))
profile.set_shape("attention_mask",
min = (batch_size, 1),
opt = (batch_size, max_seq_length//2),
max = (batch_size, max_seq_length))
profile.set_shape("token_type_ids",
min = (batch_size, 1),
opt = (batch_size, max_seq_length//2),
max = (batch_size, max_seq_length))
config.add_optimization_profile(profile)
TensorRT 将根据 min, opt 和 max 设置的值进行优化。
TensorRT Engine 加载
加载 TensorRT Engine 模型至少需要用到 logger, runtime, engine 和 context 组件，其中
- runtime 负责反序列化模型文件生成 engine 对象，
- engine 负责创建 context 对象，
- context 负责执行推理过程，并维护推理过程中的上下文，包括输入输出数据。
with open(engine_path, "rb") as f:
engineString = f.read()
logger = trt.Logger(trt.Logger.VERBOSE)
runtime = trt.Runtime(logger)
engine = runtime.deserialize_cuda_engine(engineString)
context = engine.create_execution_context()
TensorRT 推理过程
TensorRT 的推理过程可以分为四个阶段，第一阶段是设置输入参数，第二个阶段是为输出分配内存，第三个阶段是执行推理，第四个阶段是获取输出结果。
其中设置输入参数涉及到数据从 CPU 到 GPU 的拷贝过程，需要引入 cuda 模块，需要使用 pip 或 conda 安装 cuda-python 包，下面给出示例并逐行解释
import tensorrt as trt
import numpy as np
from cuda import cudart
## input 是一个字典，包含模型的输入张量名称和对应的数据
input = {"input": np.random.randn(1, 3, 224, 224)}
## 用于保存 cuda 内存指针的容器，用于后续释放内存
input_buffers = {}
## num_io_tensors 表示模型输入输出张量的数量
for i in range(engine.num_io_tensors):
## get_tensor_name 获取张量名称，就是在导出 ONNX 模型中定义的名称
name = engine.get_tensor_name(i)
## get_tensor_mode 获取张量类别，包括输入和输出两种，这里只处理输入张量
if engine.get_tensor_mode(name) != trt.TensorIOMode.INPUT:
continue
## 获取输入张量数据
array = input[name]
## 模型要求的输入张量数据类型，这里我们把转换成 numpy 类型
dtype = np.dtype(trt.nptype(engine.get_tensor_dtype(name)))
## 保持输入张量类型与模型要求的数据类型一致
array = array.astype(dtype)
## numpy 数组的内存布局有可能不是连续的，这里需要转换为连续的内存布局，以便使用指针拷贝
array = np.ascontiguousarray(array)
## cudaMalloc 分配 GPU 内存，返回内存指针和错误码
err, ptr = cudart.cudaMalloc(array.nbytes)
if err > 0:
raise Exception("cudaMalloc failed, error code: {}".format(err))
## 暂时保存内存指针，后续还需要释放
input_buffers[name] = ptr
## cudaMemcpy 将数据从 CPU 拷贝到 GPU，其中 array.ctypes.data 是 numpy 数组的内存指针
cudart.cudaMemcpy(ptr, array.ctypes.data, array.nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice)
## set_input_shape 设置输入张量的实际形状，对于 dynamic shape 这一步是必要的，因为动态维度在 ONNX 转换过程中被设置成了 -1，这里不设置将会报错
context.set_input_shape(name, array.shape)
## set_tensor_address 设置输入张量的内存地址
context.set_tensor_address(name, ptr)
不同于一般的框架，TensorRT 还需要手动设置输出数据的内存空间，可以选择 set_tensor_address 方法，但是我在 Bert 上使用该方法结果不正确，所以下面选择使用 OutputAllocator 来让 context 在运行完成后自动设置。
class OutputAllocator(trt.IOutputAllocator):
def __init__(self):
trt.IOutputAllocator.__init__(self)
self.buffers = {}
self.shapes = {}
def reallocate_output(self, tensor_name, memory, size, alignment):
ptr = cudart.cudaMalloc(size)[1]
self.buffers[tensor_name] = ptr
return ptr
def notify_shape(self, tensor_name, shape):
self.shapes[tensor_name] = tuple(shape)
这里我们定义的 OutputAllocator 继承自 IOutputAllocator，TensorRT 在分配输出内存时将调用 reallocate_output 方法，其中的参数：
- tensor_name 表示输出张量的名称，
- memory 表示
set_tensor_address方法设置的内存地址，这里我们没有使用。根据 TensorRT 的 API 文档[2]，提前使用set_tensor_address方法设置输出张量的内存地址，可以让 TensorRT 直接使用这块内存，除非内存大小不足才会调用reallocate_output方法重新分配内存。 - size 表示输出张量的内存大小，注意这是字节大小，不是数组大小，
- alignment 表示内存分配的对齐大小，这里我们没有使用。
另外，notify_shape 方法将在 TensorRT 计算出 output 的 shape 后调用。有了 OutputAllocator 之后，可以很方便地设置 context 的输出内存
output_allocator = OutputAllocator()
for i in range(engine.num_io_tensors):
name = engine.get_tensor_name(i)
if engine.get_tensor_mode(name) != trt.TensorIOMode.OUTPUT:
continue
context.set_output_allocator(name, output_allocator)
接下来是执行推理，随着 TensorRT 版本的迭代，它的 API 也有可能变化，目前版本(8.6.1)比较常用的方法如下
context.execute_async_v3(0)
其中的参数表示推理过程在哪个 stream 上执行，这里使用默认的 stream。
执行完成后，我们可以从 output_allocator 中获取输出张量的内存地址，注意这是在 GPU 上的地址，需要拷贝到 CPU 上。
output = {}
for name in output_allocator.buffers.keys():
ptr = output_allocator.buffers[name]
shape = output_allocator.shapes[name]
dtype = np.dtype(trt.nptype(engine.get_tensor_dtype(name)))
nbytes = np.prod(shape) * dtype.itemsize
output_buffer = np.empty(shape, dtype = dtype)
cudart.cudaMemcpy(output_buffer.ctypes.data, ptr, nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost)
output[name] = output_buffer
最后不要忘记释放 GPU 内存
for name in input_buffers.keys():
ptr = input_buffers[name]
cudart.cudaFree(ptr)
for name in output_allocator.buffers.keys():
ptr = output_allocator.buffers[name]
cudart.cudaFree(ptr)
总结
我们在这一节内容中介绍了 TensorRT 的最基础使用方法，包括 PyTorch 模型到 ONNX 转换，ONNX 到 Engine 转换，Engine 加载、推理过程等。即便没有使用任何优化手段，TensorRT 也能带来相较于 PyTorch 模型大致 5 倍的推理性能提升。在后续的内容中，我们将介绍如何进一步的优化模型。