中继TensorRT集成

2024-02-23 10:05:10

中继TensorRT集成

介绍

NVIDIA TensorRT是用于优化深度学习推理的库。这种集成将使尽可能多的算子从Relay转移到TensorRT，从而无需调整调度，即可在NVIDIA GPU上提高性能。

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

安装TensorRT

要下载TensorRT，需要创建一个NVIDIA Developer程序帐户。请参阅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

TensorRT在TVM中的集成有两个单独的构建标记。这些标志还可以启用交叉编译：USE_TENSORRT_CODEGEN = ON在主机上构建具有TensorRT支持的模块，而USE_TENSORRT_RUNTIME = ON使边缘设备上的TVM运行时runtime执行TensorRT模块。如果要编译，也要启用具有相同TVM构建的模型，则应同时启用。

USE_TENSORRT_CODEGEN = ON / OFF-此标志将启用编译TensorRT模块，该模块不需要任何TensorRT库。
USE_TENSORRT_RUNTIME = ON / OFF / TensorRT路径-此标志将启用TensorRT运行时runtime模块。针对已安装的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, config = partition_for_tensorrt(mod, params)

使用partition_for_tensorrt返回的新模块和配置来构建Relay图。目标必须始终是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)

导出模块。

lib.export_library('compiled.so')

加载模块并在目标计算机上运行推理，必须在USE_TENSORRT_RUNTIME启用后对其进行构建。由于必须构建TensorRT引擎，因此第一次运行会花费更长的时间。

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)

分区和编译设置

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

version-TensorRT版本以（major, minor, patch）元组为目标。如果使用USE_TENSORRT_RUNTIME = ON编译TVM，则将改用链接的TensorRT版本。版本影响哪些算子分区到TensorRT。
use_implicit_batch-使用TensorRT隐式批处理模式（默认为true）。设置为false将启用显式批处理模式，扩大支持算子范围，使其包括修改批处理维度的算子，但可能会降低某些模型的性能。
remove_no_mac_subgraphs-启发式改进性能。如果没有任何乘累加运算，则删除已为TensorRT分区的子图。删除的子图将通过TVM的标准编译。
max_workspace_size-允许每个子图用于TensorRT引擎创建的工作空间大小的字节数。有关更多信息，请参见TensorRT文档。可以在运行时runtime覆盖。

运行时runtime设置

可以在运行时runtime使用环境变量配置一些其他选项。

FP16自动转换-TVM_TENSORRT_USE_FP16=1可以设置环境变量，将模型的TensorRT组件自动转换为16位浮点精度。可以大大提高性能，但可能会导致模型精度略有下降。
缓存TensorRT引擎-在首次推理期间，运行时runtime将调用TensorRT API来构建引擎。这可能很耗时，因此可以设置TVM_TENSORRT_CACHE_DIR指向一个目录来将这些内置引擎保存到磁盘上。下次加载模型并给其提供相同目录时，运行时runtime将加载已构建的引擎，以避免长时间的预热。每个模型都需要一个唯一的目录。
TensorRT具有用于配置模型中每个图层可以使用的最大缓存cache空间的参数。通常最好使用不会导致内存不足的最大值。可以TVM_TENSORRT_MAX_WORKSPACE_SIZE通过指定要使用的字节大小，指定工作区大小来覆盖此设置。

Operator support支持

增加一个新的算子

为了增加对新算子的支持，需要对以下文件进行一系列更改：

src / runtime / contrib / tensorrt / tensorrt_ops.cc创建一个新的op转换器类来实现该TensorRTOpConverter接口。必须实现构造函数以指定有多少输入以及它们是张量还是权重。还必须实现该 Convert方法来执行转换。使用参数的输入，属性和网络来添加新的TensorRT层，部署输出来完成的。可以使用现有的转换器为例。最后，在GetOpConverters()映射图上注册新的算子标签。
python / relay / op / contrib / tensorrt.py，此文件包含TensorRT的标记规则。确定支持哪些算子及其属性。必须为中继算子注册一个注释函数，并通过检查属性返回true或false，来指定转换器支持哪些属性。
tests / python / contrib / test_tensorrt.py为给定的算子添加单元测试。

码农公寓

介绍