Mxnet使用TensorRT加速模型--Mxnet官方例子

官方示例鏈接  https://mxnet.apache.org/api/python/docs/tutorials/performance/backend/tensorrt/tensorrt.html
  Optimizing Deep Learning Computation Graphs with TensorRT
  本文檔代碼來源於mxnet官方網站翻譯，添加自己運行情況。NVIDIA的TensorRT是用於加速網絡推理的深度學習庫，Mxnet的1.3.0版本開始試驗性質的支持TensorRT，這意味着Mxnet的用戶現在可以用加速庫高效的運行他們的網絡。

安裝TensorRT
  目前Mxnet集成的TensorRT支持的環境為

Ubuntu16.04
CUDA9.0或9.2
Pascal或更新架構的顯卡
下載並安裝TensorRT庫
  當這些先決條件滿足並保證更新后，可以通過pip安裝一個支持TensorRT的Mxnet編譯版本。安裝命令為：
  CUDA9.0環境
  pip install mxnet-tensorrt-cu90
  CUDA 9.2環境
  pip install mxnet-tensorrt-cu92
  非Ubuntu 16.04的系統或者只使用docker鏡像
  nvidia-docker run -ti mxnet/tensorrt python

模型加載與轉換
import os
import mxnet as mx
from mxnet.gluon.model_zoo import vision
import time

# 轉換模型
batch_shape = (1, 3, 224, 224)
resnet18 = vision.resnet18_v2(pretrained=True)
resnet18.hybridize()
resnet18.forward(mx.nd.zeros(batch_shape))
resnet18.export('resnet18_v2')

#加載新模型
sym, arg_params, aux_params = mx.model.load_checkpoint('resnet18_v2', 0)

# Create sample input
input = mx.nd.zeros(batch_shape)

  TensorRT是一個僅用於推理的庫，官方示例中使用Resnet18作為比較TensorRT的基准模型，Resnet18是一個計算復雜度較高的網絡結構，常用於計算機視覺任務和深度學習庫性能基准測試。Mxnet的Gluon模塊可以從Gluon Model Zoo下載Resnet18模型，由於當前版本的TensorRT只支持symbolic的模型，因此需要對Gluon模塊下載的模型進行hybridize、export，官方說后續可能考慮更新到直接加速Gluon模型。 

模型直接運行性能
# Execute with MXNet
os.environ['MXNET_USE_TENSORRT'] = '0'
executor = sym.simple_bind(ctx=mx.gpu(0), data=batch_shape, grad_req='null', force_rebind=True)
executor.copy_params_from(arg_params, aux_params)

# Warmup
print('Warming up MXNet')
for i in range(0, 10):
    y_gen = executor.forward(is_train=False, data=input)
    y_gen[0].wait_to_read()

# Timing
print('Starting MXNet timed run')
start = time.process_time()
for i in range(0, 10000):
    y_gen = executor.forward(is_train=False, data=input)
    y_gen[0].wait_to_read()
end = time.time()
print(time.process_time() - start)

  為了提高測試的准確性，例子中在測試前執行了一定的推理次數對模型進行warmup，保證推理中並未實際使用的lazy operations在測試前已經完成，最終在TitanV上跑出了33.73s。

模型經TensorRT加速后運行性能
# Execute with TensorRT
print('Building TensorRT engine')
os.environ['MXNET_USE_TENSORRT'] = '1'
arg_params.update(aux_params)
all_params = dict([(k, v.as_in_context(mx.gpu(0))) for k, v in arg_params.items()])
executor = mx.contrib.tensorrt.tensorrt_bind(sym, ctx=mx.gpu(0), all_params=all_params,
                                             data=batch_shape, grad_req='null', force_rebind=True)

#Warmup
print('Warming up TensorRT')
for i in range(0, 10):
    y_gen = executor.forward(is_train=False, data=input)
    y_gen[0].wait_to_read()

# Timing
print('Starting TensorRT timed run')
start = time.process_time()
for i in range(0, 10000):
    y_gen = executor.forward(is_train=False, data=input)
    y_gen[0].wait_to_read()
end = time.time()
print(time.process_time() - start)

  當前mxnet的TensorRT還處於實驗性質，因此使用的是contrib中的tensorrt_bind接口，參數初始化時候也是按照字典加載進去的。后續TensorRT將逐步改進，tensorrt_bind將並取消並取得和simple_bind相同的接口風格。
  最終經TensorRT加速后模型的運行時間為18.99s，在本例中的加速原理來源於operator fusion。

Operators and Subgraph Fusion
  這些出現的優化是以算子和子圖的融合為基礎發生的，神經網絡可以被看做是一堆算子的計算圖。算子可以執行各種計算，但在tensor上主要執行的是簡單的數學運算或線性代數。這些算子融合在一起，在一個CUDA大核上使用通常更為高效。

Today’s deep learning models perform tens of thousands of operations on GPU. The input and output of each GPU kernel has to be stored in the global memory, but read and write on global memory is much slower than on on-chip register. When some special kernels executed in sequence share some data, performance and memory locality can be improved by fusing these kernels into a single, larger one, operating on on-chip register instead of global memory.

融合子圖到一個CUDA大核上進行計算，可以避免使用大量全局變量導致的讀寫太慢問題。

  mxnet_tensorrt能夠掃描整個計算圖，並找到可以通過TensorRT優化的子圖。因此在構造mxnet計算圖時會查找是否有支持TensorRT操作符的子圖，對可以支持的用TensorRT節點替換它們，在運行時對達到這個TensorRT節點運行TensorRT自己庫中的子圖實現，這些TensorRT子圖可以將大量算子融合到一個CUDA內核進行計算。
  Mxnet負責輸入數據到相應的TensorRT節點並獲取輸出，同時刪除掉初始化期間重復的權重參數，如只用於TensorRT部分的權重。
  Resnet在TensorRT優化前后的圖結構為：


其它問題
  1、使用例子后，實際時間沒區別，這點目前還沒找到原因，論壇上也有人在問，后續看怎么發展吧
tensorrt使用與不使用沒有區別

  2、論壇上關於TensorRT的帖子還有個提到了內存占用的問題，目前還沒實現到這一步，后面遇到了再來看看
int8模型量化內存占用極大