当前位置:首页 > Web开发 > 正文

Mxnet负责输入数据到相应的TensorRT节点并获取输出

2024-03-31 Web开发

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

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

温馨提示: 本文由Jm博客推荐,转载请保留链接: https://www.jmwww.net/file/web/32472.html