Relay張量集成

Introduction

NVIDIA TensorRT是一個(gè)用于優(yōu)化深度學(xué)習(xí)推理的庫。這種集成將盡可能多地減輕從中繼到TensorRT的算子，在NVIDIA GPU上提供性能提升，而無需調(diào)整計(jì)劃。

本文將演示如何安裝TensorRT并在啟用TensorRT BYOC和運(yùn)行時(shí)的情況下構(gòu)建TVM。將提供使用TensorRT編譯和運(yùn)行ResNet-18模型的示例代碼，以及如何配置編譯和運(yùn)行時(shí)設(shè)置。最后，記錄支持的運(yùn)算符以及如何擴(kuò)展集成以支持其他運(yùn)算符。

Installing TensorRT

為了下載TensorRT，需要?jiǎng)?chuàng)建一個(gè)NVIDIA開發(fā)人員程序帳戶。有關(guān)更多信息，請參閱NVIDIA的文檔：https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html。如果有Jetson設(shè)備，如TX1、TX2、Xavier或Nano，TensorRT將通過JetPack SDK安裝在設(shè)備上。

安裝TensorRT有兩種方法：

通過deb或rpm軟件包安裝系統(tǒng)。

Tar文件安裝。

對于tar文件安裝方法，必須提供提取的tar歸檔文件的路徑，才能USE_TENSORRT_RUNTIME=/path/to/TensorRT。對于系統(tǒng)安裝方法，USE_TENSORRT_RUNTIME=ON將自動(dòng)定位安裝。

Building TVM with TensorRT support

在TVM中有兩個(gè)獨(dú)立的構(gòu)建標(biāo)志用于TensorRT集成。這些標(biāo)志還支持交叉編譯：USE_TENSORRT_CODEGEN=ON還可以在主機(jī)上構(gòu)建支持TENSORRT的模塊，而USE_TENSORRT_RUNTIME=ON將使邊緣設(shè)備上的TVM運(yùn)行時(shí)執(zhí)行TENSORRT模塊。如果希望編譯和執(zhí)行具有相同TVM構(gòu)建的模型，則應(yīng)該同時(shí)啟用這兩個(gè)功能。

USE_TENSORRT_CODEGEN=ON/OFF-此標(biāo)志將允許編譯不需要任何TENSORRT庫的TENSORRT模塊。

USE_TENSORRT_RUNTIME=ON/OFF/path-to-TensorRT-此標(biāo)志將啟用TENSORRT運(yùn)行時(shí)模塊。這將根據(jù)已安裝的TensorRT庫構(gòu)建TVM。

示例設(shè)置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集成支持的所有操作都將被標(biāo)記并卸載到TensorRT。其余的操作將通過常規(guī)的TVM CUDA編譯和代碼生成。

from tvm.relay.op.contrib.tensorrt import partition_for_tensorrt

mod, config = partition_for_tensorrt(mod, params)

使用partition_for_tensorrt返回的新模塊和配置構(gòu)建中繼圖。目標(biāo)必須始終是cuda目標(biāo)。partition_for_tensorrt將自動(dòng)填充配置中所需的值，因此不需要修改它-只需將其傳遞給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’)

加載模塊并在目標(biāo)計(jì)算機(jī)上運(yùn)行推斷，目標(biāo)計(jì)算機(jī)必須在啟用USE_TENSORRT_RUNTIME運(yùn)行時(shí)的情況下構(gòu)建。第一次運(yùn)行需要更長的時(shí)間，因?yàn)門ensorRT engine必須編譯出來。

ctx = tvm.gpu(0)

loaded_lib = tvm.runtime.load_module(‘compiled.so’)

gen_module = tvm.contrib.graph_runtime.GraphModule(loaded_lib’default’)

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

gen_module.run(data=input_data)

Partitioning and Compilation
Settings

有一些選項(xiàng)可以在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隱式批處理模式（默認(rèn)為true）。設(shè)置為false將啟用顯式批處理模式，該模式將擴(kuò)展支持的運(yùn)算符以包括那些修改batch dimension的運(yùn)算符，但可能會(huì)降低某些模型的性能。

remove_no_mac_subgraphs-一種改進(jìn)性能的啟發(fā)式方法。刪除已為TensorRT分區(qū)的子圖（如果它們沒有任何乘法累加運(yùn)算）。刪除的子圖將經(jīng)過TVM的標(biāo)準(zhǔn)編譯。

max_workspace_size-允許每個(gè)子圖用于創(chuàng)建TensorRT引擎的工作區(qū)大小字節(jié)數(shù)。有關(guān)更多信息，請參見TensorRT文檔。可以在runtime重寫。

Runtime Settings

有一些附加選項(xiàng)可以在runtime使用環(huán)境變量進(jìn)行配置。

Automatic FP16 Conversion-可以設(shè)置環(huán)境變量TVM_TENSORRT_USE_FP16=1，以自動(dòng)將模型的TENSORRT組件轉(zhuǎn)換為16位浮點(diǎn)精度。這可以極大地提高性能，但可能會(huì)在模型精度方面造成一些輕微的損失。

緩存TensorRT引擎-在第一次推斷期間，runtime將調(diào)用TensorRT API來構(gòu)建引擎。這可能很耗時(shí)，因此可以將TVM_TENSORRT_CACHE_DIR設(shè)置為指向?qū)⑦@些內(nèi)置引擎保存到磁盤上的目錄。下次加載模型并給它相同的目錄時(shí)，runtime將加載已經(jīng)構(gòu)建的引擎，以避免長時(shí)間的預(yù)熱。每個(gè)模型都需要一個(gè)唯一的目錄。

TensorRT有一個(gè)參數(shù)來配置模型中每個(gè)層可以使用的最大暫存空間量。通常最好使用不會(huì)導(dǎo)致內(nèi)存不足的最大值。可以使用TVM_TENSORRT_MAX_WORKSPACE_SIZE來覆蓋此設(shè)置，方法是指定要使用的工作區(qū)大小（以字節(jié)為單位）。

Operator support

Adding a new operator

要添加對新算子的支持，需要對一系列文件進(jìn)行更改：

src/runtime/contrib/tensorrt/tensorrt_ops.cc抄送創(chuàng)建一個(gè)實(shí)現(xiàn)TensorRTOpConverter接口的新op converter類。必須實(shí)現(xiàn)構(gòu)造函數(shù)來指定有多少個(gè)輸入以及它們是張量還是權(quán)重。必須實(shí)現(xiàn)Convert方法才能執(zhí)行轉(zhuǎn)換。這是通過使用params中的輸入、屬性和網(wǎng)絡(luò)來添加新的TensorRT層并推送層輸出來完成的。

可以使用現(xiàn)有的轉(zhuǎn)換器作為示例。

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

python/relay/op/contrib/tensorrt.py此文件包含TensorRT的注釋規(guī)則。決定支持哪些運(yùn)算符及其屬性。必須為中繼運(yùn)算符注冊注釋函數(shù)，并通過檢查屬性是否返回true或false來指定轉(zhuǎn)換器支持哪些屬性。

tests/python/contrib/test_tensorrt.py為給定的運(yùn)算符添加單元測試。

總結(jié)

以上是生活随笔為你收集整理的Relay张量集成的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。