Mxnet TensorRT

Optimizing Deep Learning Computation Graphs with TensorRT

CUDA9.0或9.2
Pascal或更新架构的显卡
下载并安装TensorRT库
  当这些先决条件满足并保证更新后，可以通过pip安装一个支持TensorRT的Mxnet编译版本。安装命令为：
  CUDA9.0环境
  pip install mxnet-tensorrt-cu90
  CUDA 9.2环境
  pip install mxnet-tensorrt-cu100
  非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)

我这运行是报错的

AttributeError: module 'mxnet.contrib.tensorrt' has no attribute 'tensorrt_bind'

  当前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部分的权重。