Relay张量集成

Relay张量集成

Introduction

NVIDIA TensorRT是一个用于优化深度学习推理的库。这种集成将尽可能多地减轻从中继到TensorRT的算子，在NVIDIA GPU上提供性能提升，而无需调整计划。

本文将演示如何安装TensorRT并在启用TensorRT BYOC和运行时的情况下构建TVM。将提供使用TensorRT编译和运行ResNet-18模型的示例代码，以及如何配置编译和运行时设置。最后，记录支持的运算符以及如何扩展集成以支持其他运算符。

Installing 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将自动定位安装。

Building TVM with TensorRT support

在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)

Build and Deploy ResNet-18 with TensorRT

Create a Relay graph from a MXNet ResNet-18 model.

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, config = partition_for_tensorrt(mod, params)

使用partition_for_tensorrt返回的新模块和配置构建中继图。目标必须始终是cuda目标。partition_for_tensorrt将自动填充配置中所需的值，因此不需要修改它-只需将其传递给PassContext，这样就可以在编译期间读取这些值。

target = "cuda"

with tvm.transform.PassContext(opt_level=3, config={'relay.ext.tensorrt.options': config}):

    lib = relay.build(mod, target=target, params=params)

Export the module.

lib.export_library('compiled.so')

加载模块并在目标计算机上运行推断，目标计算机必须在启用USE_TENSORRT_RUNTIME运行时的情况下构建。第一次运行需要更长的时间，因为TensorRT engine必须编译出来。

ctx = tvm.gpu(0)

loaded_lib = tvm.runtime.load_module('compiled.so')

gen_module = tvm.contrib.graph_runtime.GraphModule(loaded_lib['default'](ctx))

input_data = np.random.uniform(0, 1, input_shape).astype(dtype)

gen_module.run(data=input_data)

Partitioning and Compilation Settings

有一些选项可以在partition_for_tensorrt配置。

version-TensorRT version to target as tuple of (major, minor, patch)。如果TVM是使用USE_TENSORRT_RUNTIME=ON编译的，则将改用链接的TENSORRT版本。算子分解到TensorRT，将影响版本。

use_implicit_batch-使用TensorRT隐式批处理模式（默认为true）。设置为false将启用显式批处理模式，该模式将扩展支持的运算符以包括那些修改batch dimension的运算符，但可能会降低某些模型的性能。

remove_no_mac_subgraphs-一种改进性能的启发式方法。删除已为TensorRT分区的子图（如果它们没有任何乘法累加运算）。删除的子图将经过TVM的标准编译。

max_workspace_size-允许每个子图用于创建TensorRT引擎的工作区大小字节数。有关更多信息，请参见TensorRT文档。可以在runtime重写。

Runtime Settings

有一些附加选项可以在runtime使用环境变量进行配置。

Automatic FP16 Conversion-可以设置环境变量TVM_TENSORRT_USE_FP16=1，以自动将模型的TENSORRT组件转换为16位浮点精度。这可以极大地提高性能，但可能会在模型精度方面造成一些轻微的损失。

缓存TensorRT引擎-在第一次推断期间，runtime将调用TensorRT API来构建引擎。这可能很耗时，因此可以将TVM_TENSORRT_CACHE_DIR设置为指向将这些内置引擎保存到磁盘上的目录。下次加载模型并给它相同的目录时，runtime将加载已经构建的引擎，以避免长时间的预热。每个模型都需要一个唯一的目录。

TensorRT有一个参数来配置模型中每个层可以使用的最大暂存空间量。通常最好使用不会导致内存不足的最大值。可以使用TVM_TENSORRT_MAX_WORKSPACE_SIZE来覆盖此设置，方法是指定要使用的工作区大小（以字节为单位）。

Operator support

Adding a new operator

要添加对新算子的支持，需要对一系列文件进行更改：

src/runtime/contrib/tensorrt/tensorrt_ops.cc抄送创建一个实现TensorRTOpConverter接口的新op converter类。必须实现构造函数来指定有多少个输入以及它们是张量还是权重。必须实现Convert方法才能执行转换。这是通过使用params中的输入、属性和网络来添加新的TensorRT层并推送层输出来完成的。

可以使用现有的转换器作为示例。

最后，在GetOpConverters（）映射中注册新的op conventer。

python/relay/op/contrib/tensorrt.py此文件包含TensorRT的注释规则。决定支持哪些运算符及其属性。必须为中继运算符注册注释函数，并通过检查属性是否返回true或false来指定转换器支持哪些属性。

tests/python/contrib/test_tensorrt.py为给定的运算符添加单元测试。