Relay TensorRT 集成#
Author: Trevor Morris
简介#
NVIDIA TensorRT 是用于优化深度学习推理的库。此集成将尽可能多地将算子从 Relay 卸载到 TensorRT，从而在 NVIDIA GPU 上提供性能提升，而无需调整调度。
本指南将演示如何安装 TensorRT 并构建启用 TensorRT BYOC 和运行时的 TVM。它还将提供使用 TensorRT 编译和运行 ResNet-18 模型的示例代码，以及如何配置编译和运行时设置。最后，将记录支持的算子以及如何扩展集成以支持其他算子。
安装 TensorRT#
要下载 TensorRT，您需要创建 NVIDIA 开发者计划账户。请参阅 NVIDIA 的文档以获取更多信息：https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html。如果您拥有 Jetson 设备（例如 TX1、TX2、Xavier 或 Nano），TensorRT 将通过 JetPack SDK 预装在设备上。
有两种安装 TensorRT 的方法：
通过 deb 或 rpm 包进行系统安装。
通过 tar 文件安装。
使用 tar 文件安装方法时，您必须提供解压后的 tar 存档路径，例如 USE_TENSORRT_RUNTIME=/path/to/TensorRT。使用系统安装方法时，USE_TENSORRT_RUNTIME=ON 将自动定位您的安装。
构建支持 TensorRT 的 TVM#
TVM 中的 TensorRT 集成有两个独立的构建标志。这些标志还支持交叉编译：USE_TENSORRT_CODEGEN=ON 允许您在主机上构建支持 TensorRT 的模块，而 USE_TENSORRT_RUNTIME=ON 将启用边缘设备上的 TVM 运行时以执行 TensorRT 模块。如果您想使用相同的 TVM 构建编译并执行模型，则应同时启用这两个标志。
USE_TENSORRT_CODEGEN=ON/OFF - 此标志将启用编译 TensorRT 模块，该模块不需要任何 TensorRT 库。
USE_TENSORRT_RUNTIME=ON/OFF/path-to-TensorRT - 此标志将启用 TensorRT 运行时模块。这将根据已安装的 TensorRT 库构建 TVM。
config.cmake 文件中的示例设置：
set(USE_TENSORRT_CODEGEN ON)
set(USE_TENSORRT_RUNTIME /home/ubuntu/TensorRT-7.0.0.11)
使用 TensorRT 构建和部署 ResNet-18#
从 MXNet ResNet-18 模型创建 Relay 图。
import tvm
from tvm import relay
import mxnet
from mxnet.gluon.model_zoo.vision import get_model
dtype = "float32"
input_shape = (1, 3, 224, 224)
block = get_model('resnet18_v1', pretrained=True)
mod, params = relay.frontend.from_mxnet(block, shape={'data': input_shape}, dtype=dtype)
为 TensorRT 注解和分区图。TensorRT 集成支持的所有算子将被标记并卸载到 TensorRT。其余算子将通过常规的 TVM CUDA 编译和代码生成。
from tvm.relay.op.contrib.tensorrt import partition_for_tensorrt
mod = partition_for_tensorrt(mod, params)
使用 partition_for_tensorrt 返回的新模块和配置构建 Relay 图。目标必须始终是 cuda 目标。partition_for_tensorrt 将自动填充配置中的所需值，因此无需修改它——只需将其传递给 PassContext，以便在编译期间读取这些值。
target = "cuda"
with tvm.transform.PassContext(opt_level=3):
lib = relay.build(mod, target=target, params=params)
导出模块。
lib.export_library('compiled.so')
在目标机器上加载模块并运行推理，目标机器必须启用 USE_TENSORRT_RUNTIME 构建。第一次运行会花费较长时间，因为需要构建 TensorRT 引擎。
dev = tvm.cuda(0)
loaded_lib = tvm.runtime.load_module('compiled.so')
gen_module = tvm.contrib.graph_executor.GraphModule(loaded_lib['default'](dev))
input_data = np.random.uniform(0, 1, input_shape).astype(dtype)
gen_module.run(data=input_data)
分区和编译设置#
可以在 partition_for_tensorrt 中配置一些选项。
version- 目标 TensorRT 版本，格式为 (major, minor, patch) 元组。如果 TVM 使用 USE_TENSORRT_RUNTIME=ON 编译，则将使用链接的 TensorRT 版本。版本将影响哪些算子可以分区到 TensorRT。use_implicit_batch- 使用 TensorRT 隐式批处理模式（默认为 true）。设置为 false 将启用显式批处理模式，这将扩展支持的算子以包括修改批处理维度的算子，但可能会降低某些模型的性能。remove_no_mac_subgraphs- 一种提高性能的启发式方法。如果子图没有乘加操作，则删除为 TensorRT 分区的子图。删除的子图将通过 TVM 的标准编译进行处理。max_workspace_size- 允许每个子图用于 TensorRT 引擎创建的工作空间大小（以字节为单位）。有关更多信息，请参阅 TensorRT 文档。可以在运行时覆盖。
运行时设置#
可以使用环境变量在运行时配置一些其他选项。
自动 FP16 转换 - 可以设置环境变量
TVM_TENSORRT_USE_FP16=1以自动将模型的 TensorRT 组件转换为 16 位浮点精度。这可以大大提高性能，但可能会导致模型精度略有下降。缓存 TensorRT 引擎 - 在第一次推理期间，运行时将调用 TensorRT API 来构建引擎。这可能很耗时，因此您可以设置
TVM_TENSORRT_CACHE_DIR指向一个目录，以将这些构建的引擎保存到磁盘上。下次加载模型并提供相同的目录时，运行时将加载已构建的引擎以避免长时间预热。每个模型都需要一个唯一的目录。TensorRT 有一个参数用于配置模型中每层可以使用的最大暂存空间量。通常最好使用不会导致内存不足的最大值。您可以使用
TVM_TENSORRT_MAX_WORKSPACE_SIZE通过指定要使用的工作空间大小（以字节为单位）来覆盖此设置。对于包含动态批处理维度的模型，可以使用变量
TVM_TENSORRT_MULTI_ENGINE来确定如何在运行时创建 TensorRT 引擎。默认模式TVM_TENSORRT_MULTI_ENGINE=0将一次仅在内存中维护一个引擎。如果遇到具有更高批处理大小的输入，引擎将使用新的 max_batch_size 设置重建。该引擎将兼容从 1 到 max_batch_size 的所有批处理大小。此模式减少了运行时使用的内存量。第二种模式TVM_TENSORRT_MULTI_ENGINE=1将为遇到的每个批处理大小构建一个唯一的 TensorRT 引擎。这将提供更高的性能，但会消耗更多内存。
支持的算子#
Relay 节点 |
标记 |
|---|---|
nn.relu |
|
sigmoid |
|
tanh |
|
nn.batch_norm |
|
nn.layer_norm |
|
nn.softmax |
|
nn.conv1d |
|
nn.conv2d |
|
nn.dense |
|
nn.bias_add |
|
add |
|
subtract |
|
multiply |
|
divide |
|
power |
|
maximum |
|
minimum |
|
nn.max_pool2d |
|
nn.avg_pool2d |
|
nn.global_max_pool2d |
|
nn.global_avg_pool2d |
|
exp |
|
log |
|
sqrt |
|
abs |
|
negative |
|
nn.batch_flatten |
|
expand_dims |
|
squeeze |
|
concatenate |
|
nn.conv2d_transpose |
|
transpose |
|
layout_transform |
|
reshape |
|
nn.pad |
|
sum |
|
prod |
|
max |
|
min |
|
mean |
|
nn.adaptive_max_pool2d |
|
nn.adaptive_avg_pool2d |
|
nn.batch_matmul |
|
clip |
Requires TensorRT 5.1.5 or greater |
nn.leaky_relu |
Requires TensorRT 5.1.5 or greater |
sin |
Requires TensorRT 5.1.5 or greater |
cos |
Requires TensorRT 5.1.5 or greater |
atan |
Requires TensorRT 5.1.5 or greater |
ceil |
Requires TensorRT 5.1.5 or greater |
floor |
Requires TensorRT 5.1.5 or greater |
split |
Requires TensorRT 5.1.5 or greater |
strided_slice |
Requires TensorRT 5.1.5 or greater |
nn.conv3d |
Requires TensorRT 6.0.1 or greater |
nn.max_pool3d |
Requires TensorRT 6.0.1 or greater |
nn.avg_pool3d |
Requires TensorRT 6.0.1 or greater |
nn.conv3d_transpose |
Requires TensorRT 6.0.1 or greater |
erf |
Requires TensorRT 7.0.0 or greater |
Adding a new operator#
要添加对新算子的支持，需要修改一系列文件：
src/runtime/contrib/tensorrt/tensorrt_ops.cc 创建实现
TensorRTOpConverter接口的新算子转换器类。您必须实现构造函数以指定有多少输入以及它们是张量还是权重。您还必须实现Convert方法以执行转换。这是通过使用来自 params 的输入、属性和网络来添加新的 TensorRT 层并推送层输出来完成的。您可以使用现有的转换器作为示例。最后，在GetOpConverters()映射中注册您的新算子转换器。python/relay/op/contrib/tensorrt.py 此文件包含 TensorRT 的注释规则。这些规则决定了支持的算子及其属性。您必须为 relay 算子注册注解函数，并通过检查属性返回 true 或 false 来指定您的转换器支持哪些属性。
tests/python/contrib/test_tensorrt.py 为给定的算子添加单元测试。