中继TensorRT集成
NVIDIA TensorRT是用于优化深度学习推理的库。这种集成将使尽可能多的算子从Relay转移到TensorRT,从而无需调整调度,即可在NVIDIA GPU上提高性能。
本文将演示如何安装TensorRT,并在启用TensorRT BYOC和运行时runtime的情况下构建TVM。将提供示例代码,使用TensorRT编译和运行ResNet-18模型,以及如何配置编译和运行时runtime设置。最后,记录支持的算子,以及如何扩展集成,以支持其它算子。
要下载TensorRT,需要创建一个NVIDIA Developer程序帐户。请参阅NVIDIA文档以获取更多信息:https : //docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html。如果有Jetson设备,例如TX1,TX2,Xavier或Nano,TensorRT将已经通过JetPack SDK安装在设备上。
有两种安装TensorRT的方法:
使用tar文件安装方法,必须将提取的tar存储文件的路径提供给USE_TENSORRT_RUNTIME = / path / to / TensorRT。使用系统安装方法,USE_TENSORRT_RUNTIME = ON将自动安装。
TensorRT在TVM中的集成有两个单独的构建标记。这些标志还可以启用交叉编译:USE_TENSORRT_CODEGEN = ON在主机上构建具有TensorRT支持的模块,而USE_TENSORRT_RUNTIME = ON使边缘设备上的TVM运行时runtime执行TensorRT模块。如果要编译,也要启用具有相同TVM构建的模型,则应同时启用。
config.cmake文件中的示例设置:
set(USE_TENSORRT_CODEGEN ON)
set(USE_TENSORRT_RUNTIME /home/ubuntu/TensorRT-7.0.0.11)
从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中配置一些选项。
可以在运行时runtime使用环境变量配置一些其他选项。
增加一个新的算子
为了增加对新算子的支持,需要对以下文件进行一系列更改:
原文:https://www.cnblogs.com/wujianming-110117/p/14176561.html