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为给定的运算符添加单元测试。