tensorrt支持的层_如何利用TensorRT加速GPU上深度学习模型推理

背景

58人工智能平台WPAI是58同城一站式算法开发平台，包括机器学习和深度学习两部分。其中深度学习平台支持GPU/CPU上深度学习模型的代码调试、离线训练和在线推理等功能。详细介绍可参考《58人工智能平台WPAI设计与实现》。

深度学习模型在线推理实现模型的生产环境部署，WPAI支持TensorFlow、PyTorch以及自定义模型的在线推理(详细介绍可参考《58深度学习在线预测服务的设计与实现》)。对于计算量较大的模型CPU上推理耗时较高，需要使用GPU资源，而GPU资源十分昂贵且资源有限，如何优化深度学习模型在GPU上的推理性能节省GPU资源成为十分有意义的事情。

TensorRT简介

TensorRT是NVIDIA 推出的一款基于CUDA和cuDNN的神经网络推理加速引擎，其核心是一个C++库，能够在NVIDIA的GPU上进行模型推理性能的优化，以下简称TRT。TRT只支持推理不支持训练，兼容当前主流的TensorFlow、Caffe2、PyTorch、Mxnet等主流深度学习框架。

图1 TensorRT功能说明

TRT加速GPU推理流程如图2所示，训练好的模型网络通过TRT进行计算图优化、精度转换与缩放、自动选取最优CUDA kernel、显存优化等操作得到优化后的执行引擎，提高推理性能。大部分深度学习框架在训练神经网络时网络中的张量都是32位浮点数的精度&#x