从零开始 TensorRT（2）Python 篇：原生 API 构建网络

前言

学习资料：
TensorRT 源码示例
官方文档：Working With TensorRT Using The Python API
官方文档：TensorRT Python API
官方文档：CUDA Python
B站视频：TensorRT 教程 | 基于 8.6.1 版本
B站视频配套代码 cookbook

  在 Python 篇中会仿照视频教程中的内容，按以下顺序介绍：

（1）TensorRT 原生 API 构建网络：示例中使用 TensorRT 的 API 搭建网络并进行推理，由于网络过于简单（只有一层且输入与输出相同），会存在如何实际搭建一个网络模型、搭建网络后如何加载参数等疑问。但作为第一个示例，主要的作用是熟悉构建网络和执行推理的流程，这部分代码比较模板化，会经常用到。

（2）解析 ONNX 模型：示例中会把一个 Pytorch 模型导出为 ONNX 格式，然后用 TensorRT 读取并生成引擎进行推理，这也是目前最常用的方法。

（3）PyTorch 框架内 TensorRT 接口：使用 Pytorch 的 API 生成 TensorRT 引擎并推理，此示例最为简单。

PyCharm 环境问题

  初学 TensorRT 时遇到了这样的问题，拿到的示例代码在终端运行时一切正常，到 PyCharm 中明明使用了相同的 conda 环境，却无法正常运行。

Bug1：ImportError: libnvinfer.so.8: cannot open shared object file: No such file or directory
Bug2：TypeError: pybind11::init(): factory function returned nullptr
Bug3：[TRT] [E] 6: [libLoader.cpp::Impl::293] Error Code 6: Internal Error (Unable to load library: libnvinfer_builder_resource.so.8.6.1: libnvinfer_builder_resource.so.8.6.1: cannot open shared object file: No such file or directory)

最终解决方法

PyCharm 顶部菜单栏 → 运行 → 编辑配置 → 环境变量一栏（没有可在修改选项中勾选） → 添加 LD_LIBRARY_PATH=xxx/TensorRT-8.6.1.6/lib

修改后只对运行当前的 Python 文件有效，也可以在编辑配置模板中一同修改，这样新的运行文件也会用同样的配置。

解决过程

  首先碰到 Bug1，查到的解决办法是在 /etc/ld.so.conf 中添加 TensorRT 的 lib 路径，并刷新 sudo ldconfig。于是便碰到了 Bug2 和 Bug3，查到的解决办法是把 libnvinfer_builder_resource.so.8.6.1 复制到 /usr/lib 下。
  虽然程序能正常运行了，但这显然是个不太好的解决方案。

  修改 /etc/ld.so.conf 和 ~/.bashrc 都是添加库文件搜索路径，而又有所不同。
  ~/.bashrc 是 Bash shell 的配置文件，对于并非使用 Bash shell 启动的应用程序可能不会生效。而 /etc/ld.so.conf 是系统级别的配置，对终端、应用程序、IDE都有效。

  以下为个人做了各种测试后的一些发现：

（1）由于配置 LD_LIBRARY_PATH 会在终端生效，所以程序能正常运行，但在 PyCharm 环境并不生效，所以会出现找不到 libnvinfer.so.8 的报错，此时整个 TRT 的 lib 目录对于 PyCharm 都不是搜索路径。

（2）在配置 /etc/ld.so.conf 后，对 PyCharm 生效却依然报错，通过 ldconfig -p | grep nvinfer 查看动态库缓存可以发现，里面只有 do_not_link_against_nvinfer_builder_resource，打开 lib 目录可以发现，这个文件是 libnvinfer_builder_resource.so.8.6.1 的链接。经测试，这个链接是在配置 /etc/ld.so.conf 后才新生成的，文件名也直接说了不要链接 nvinfer_builder_resource。虽然没有探究为什么会新生成这样一个链接，但基本可以确定 TRT 并不支持这样配置环境。

（3）在配置 /etc/ld.so.conf 后，报错转变成了 Bug2 和 Bug3 其实就是找不到 libnvinfer_builder_resource.so.8.6.1，所以暴力的把它复制到 /usr/lib 下（系统默认搜索路径），是可以解决问题的。

（4）结合以上所有的发现，可以得知问题的根本是 lib 路径不在 PyCharm 环境中，使用系统级别配置却有个别文件会出现异常，这才导致了各种报错。从而可以想到，只要在 PyCharm 环境变量中加入 lib 路径，即可解决这种问题。

示例：TensorRT 原生 API 构建网络

参考源码：cookbook → 01-SimpleDemo → TensorRT8.5 → main.py

源码

  对参考源码进行了一些简化和修改：

1. 去除从文件读取引擎部分
2. 省略了错误判断
3. 网络只有一个输入张量和一个输出张量，内存和显存操作部分去除了循环操作

import numpy as np
import tensorrt as trt
from cuda import cudart

logger = trt.Logger(trt.Logger.ERROR)
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
profile = builder.create_optimization_profile()
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)

inputTensor = network.add_input("inputT0", trt.float32, [-1, -1, -1])
profile.set_shape(inputTensor.name, [1, 1, 1], [3, 4, 5], [6, 8, 10])
config.add_optimization_profile(profile)
identityLayer = network.add_identity(inputTensor)
network.mark_output(identityLayer.get_output(0))
engineString = builder.build_serialized_network(network, config)
with open("./model.plan", "wb") as f:
    f.write(engineString)


engine = trt.Runtime(logger).deserialize_cuda_engine(engineString)
nIO = engine.num_io_tensors
lTensorName = [engine.get_tensor_name(i) for i in range(nIO)]
context = engine.create_execution_context()
context.set_input_shape(lTensorName[0], [3, 4, 5])

data = np.arange(3 * 4 * 5, dtype=np.float32).reshape(3, 4, 5)
inputHost = np.ascontiguousarray(data)
outputHost = np.empty(
    context.get_tensor_shape(lTensorName[1]),
    dtype=trt.nptype(engine.get_tensor_dtype(lTensorName[1]))
)
_, inputDevice = cudart.cudaMalloc(inputHost.nbytes)
_, outputDevice = cudart.cudaMalloc(outputHost.nbytes)
context.set_tensor_address(lTensorName[0], inputDevice)
context.set_tensor_address(lTensorName[1], outputDevice)

cudart.cudaMemcpy(inputDevice, inputHost.ctypes.data, inputHost.nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice)
context.execute_async_v3(0)
cudart.cudaMemcpy(outputHost.ctypes.data, outputDevice, outputHost.nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost)

print(lTensorName[0])
print(inputHost)
print(lTensorName[1])
print(outputHost)

cudart.cudaFree(inputDevice)
cudart.cudaFree(outputDevice)

逐句分析

（1）构建期

"日志记录器 Logger"
"多个 Builder 可共享同个 Logger, 日志等级: VERBOSE, INFO, WARNING, ERROR, INTERNAL_ERROR"
"通常使用 VERBOSE, INFO 可查看优化信息"
logger = trt.Logger(trt.Logger.ERROR)
"计算图构建器 Builder"
"模型搭建的入口, 网络的内部表示和可执行程序引擎有它的成员方法生成"
builder = trt.Builder(logger)
"计算网络(计算图) Network"
"网络的主体, 搭建网络时往里添加层并标记输入输出节点"
"EXPLICIT_BATCH 模式: 显示batch维度, 可支持 Batch Normalization, Reshape, Transpose, Reduce 等 batch 维度参与的各种变换"
"IMPLICIT_BATCH 模式: 隐藏batch维度, 老版本默认模式, 新版本逐渐弃用"
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
"动态模式(Dynamic Shape Mode)下的配置器 Optimization Profile, 帮助网络优化"
"动态模式下, 输入张量形状可以在推理时决定"
profile = builder.create_optimization_profile()

"构建器的配置器 Config"
"用于设置模型参数, 例如是否开启 FP16, INT8 模式等"
config = builder.create_builder_config()
"设置可用显存(单位: Byte)"
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)

"网络输入张量, 有3个维度且每个维度的尺寸都可变"
inputTensor = network.add_input("inputT0", trt.float32, [-1, -1, -1])
"给定输入张量的最小、最常见、最大尺寸"
"profile.set_shape(tensorName, minShape, commonShape, maxShape)"
profile.set_shape(inputTensor.name, [1, 1, 1], [3, 4, 5], [6, 8, 10])
config.add_optimization_profile(profile)

"添加 identity 层(恒等函数)"
identityLayer = network.add_identity(inputTensor)
"网络输出"
network.mark_output(identityLayer.get_output(0))
"生成序列化网络"
engineString = builder.build_serialized_network(network, config)

（2）运行期

"计算图可执行程序 Engine"
engine = trt.Runtime(logger).deserialize_cuda_engine(engineString)
nIO = engine.num_io_tensors
lTensorName = [engine.get_tensor_name(i) for i in range(nIO)]
"GPU进程"
context = engine.create_execution_context()
"推理时实际输入的形状"
context.set_input_shape(lTensorName[0], [3, 4, 5])


data = np.arange(3 * 4 * 5, dtype=np.float32).reshape(3, 4, 5)
"CPU 内存申请"
inputHost = np.ascontiguousarray(data)
outputHost = np.empty(
    context.get_tensor_shape(lTensorName[1]),
    dtype=trt.nptype(engine.get_tensor_dtype(lTensorName[1]))
)
"GPU 显存申请"
_, inputDevice = cudart.cudaMalloc(inputHost.nbytes)
_, outputDevice = cudart.cudaMalloc(outputHost.nbytes)
"设置张量显存地址"
context.set_tensor_address(lTensorName[0], inputDevice)
context.set_tensor_address(lTensorName[1], outputDevice)
"输入数据从内存拷贝到显存"
cudart.cudaMemcpy(inputDevice, inputHost.ctypes.data, inputHost.nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice)
"模型推理"
context.execute_async_v3(0)
"输出数据从显存拷贝到内存"
cudart.cudaMemcpy(outputHost.ctypes.data, outputDevice, outputHost.nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost)
"释放显存"
cudart.cudaFree(inputDevice)
cudart.cudaFree(outputDevice)

（3）从文件读取
  省略构建期直接读取文件，可以节省大量时间，但是只能在当前的软硬件环境下使用

with open("./model.plan", "rb") as f:
    engineString = f.read()