Relay张量集成
Relay張量集成
Introduction
NVIDIA TensorRT是一個用于優(yōu)化深度學習推理的庫。這種集成將盡可能多地減輕從中繼到TensorRT的算子，在NVIDIA GPU上提供性能提升，而無需調(diào)整計劃。
本文將演示如何安裝TensorRT并在啟用TensorRT BYOC和運行時的情況下構建TVM。將提供使用TensorRT編譯和運行ResNet-18模型的示例代碼，以及如何配置編譯和運行時設置。最后，記錄支持的運算符以及如何擴展集成以支持其他運算符。
Installing TensorRT
為了下載TensorRT，需要創(chuàng)建一個NVIDIA開發(fā)人員程序帳戶。有關更多信息，請參閱NVIDIA的文檔：https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html。如果有Jetson設備，如TX1、TX2、Xavier或Nano，TensorRT將通過JetPack SDK安裝在設備上。
安裝TensorRT有兩種方法：
通過deb或rpm軟件包安裝系統(tǒng)。
Tar文件安裝。
對于tar文件安裝方法，必須提供提取的tar歸檔文件的路徑，才能USE_TENSORRT_RUNTIME=/path/to/TensorRT。對于系統(tǒng)安裝方法，USE_TENSORRT_RUNTIME=ON將自動定位安裝。
Building TVM with TensorRT support
在TVM中有兩個獨立的構建標志用于TensorRT集成。這些標志還支持交叉編譯：USE_TENSORRT_CODEGEN=ON還可以在主機上構建支持TENSORRT的模塊，而USE_TENSORRT_RUNTIME=ON將使邊緣設備上的TVM運行時執(zhí)行TENSORRT模塊。如果希望編譯和執(zhí)行具有相同TVM構建的模型，則應該同時啟用這兩個功能。
USE_TENSORRT_CODEGEN=ON/OFF-此標志將允許編譯不需要任何TENSORRT庫的TENSORRT模塊。
USE_TENSORRT_RUNTIME=ON/OFF/path-to-TensorRT-此標志將啟用TENSORRT運行時模塊。這將根據(jù)已安裝的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。其余的操作將通過常規(guī)的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’)
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-一種改進性能的啟發(fā)式方法。刪除已為TensorRT分區(qū)的子圖（如果它們沒有任何乘法累加運算）。刪除的子圖將經(jīng)過TVM的標準編譯。
max_workspace_size-允許每個子圖用于創(chuàng)建TensorRT引擎的工作區(qū)大小字節(jié)數(shù)。有關更多信息，請參見TensorRT文檔?？梢栽趓untime重寫。
Runtime Settings
有一些附加選項可以在runtime使用環(huán)境變量進行配置。
Automatic FP16 Conversion-可以設置環(huán)境變量TVM_TENSORRT_USE_FP16=1，以自動將模型的TENSORRT組件轉換為16位浮點精度。這可以極大地提高性能，但可能會在模型精度方面造成一些輕微的損失。
緩存TensorRT引擎-在第一次推斷期間，runtime將調(diào)用TensorRT API來構建引擎。這可能很耗時，因此可以將TVM_TENSORRT_CACHE_DIR設置為指向?qū)⑦@些內(nèi)置引擎保存到磁盤上的目錄。下次加載模型并給它相同的目錄時，runtime將加載已經(jīng)構建的引擎，以避免長時間的預熱。每個模型都需要一個唯一的目錄。
TensorRT有一個參數(shù)來配置模型中每個層可以使用的最大暫存空間量。通常最好使用不會導致內(nèi)存不足的最大值?？梢允褂肨VM_TENSORRT_MAX_WORKSPACE_SIZE來覆蓋此設置，方法是指定要使用的工作區(qū)大小（以字節(jié)為單位）。
Operator support
Adding a new operator
要添加對新算子的支持，需要對一系列文件進行更改：
src/runtime/contrib/tensorrt/tensorrt_ops.cc抄送創(chuàng)建一個實現(xiàn)TensorRTOpConverter接口的新op converter類。必須實現(xiàn)構造函數(shù)來指定有多少個輸入以及它們是張量還是權重。必須實現(xiàn)Convert方法才能執(zhí)行轉換。這是通過使用params中的輸入、屬性和網(wǎng)絡來添加新的TensorRT層并推送層輸出來完成的。
可以使用現(xiàn)有的轉換器作為示例。
最后，在GetOpConverters（）映射中注冊新的op conventer。
python/relay/op/contrib/tensorrt.py此文件包含TensorRT的注釋規(guī)則。決定支持哪些運算符及其屬性。必須為中繼運算符注冊注釋函數(shù)，并通過檢查屬性是否返回true或false來指定轉換器支持哪些屬性。
tests/python/contrib/test_tensorrt.py為給定的運算符添加單元測試。
總結
- 上一篇： NVIDIA Turing Archit
- 下一篇： TVM部署和集成Deploy and I