官方示例链接 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模型量化内存占用极大