每日一歌,分享好心情:[《爱情转移》]

最近一直在折腾TF-TRT的问题…过程 一言难尽,记录些许心得,希望你少走弯路

一 基本概念

参考官方文档:https://docs.nvidia.com/deeplearning/frameworks/tf-trt-user-guide/index.html#introduction
某乎上的一篇入门文章: https://zhuanlan.zhihu.com/p/371239130

  • Tensorflow
    TF一个开源的、端到端的(end-to-end)机器学习框架。最初是由 Google 机器智能研究组织内的 Google Brain 团队的研究人员和工程师开发的,用于进行机器学习和深度神经网络研究, 有大量用于训练、部署、服务的官方和第三方工具/平台。
  • TensorRT
    TensorRT 是NVIDIA用于高性能深度学习推理的 SDK。包含深度学习推理优化器和运行时环境,可为深度学习推理应用提供低延迟和高吞吐量。TensorRT 以NVIDIA 的并行编程模型 CUDA 为基础构建而成,可以从每个深度学习框架(例如:tensorflow)中将已训练模型导入到 TensorRT,应用优化后,TensorRT 选择平台特定的内核,更大限度提升GPU 的性能。
  • TF-TRT
    TF-TRT是 TensorFlow with TensorRT 的缩写。 可以理解为经过TensorRT优化的Tensorflow, 依然是tensorflow 哦。 优化手段一般是将TF模型中多个Tensorflow subgraph 合并成一个或多个经过优化的TensorRT计算节点(TRTEngineOp),另外也支持FP16低精度和Int8量化。
    部署经过训练的TF模型工作流程(不适用TF-TRT):
    在这里插入图片描述使用TF-TRT 进行优化是在部署模型进行推理之前需要执行的额外步骤
    在这里插入图片描述
    总结一下:
  • TF是训练和部署的大框架,TRT是一个推理库,TF-TRT是经过TRT优化的TF
  • 注意: TRT、TF-TRF是GPU强相关的,模型转换环境必须和部署环境一致。

二 模型转换

使用tensorflow提供了TF-TRT转换接口

1. 关于转换环境

需要使用启用了TensorRT的Tensorflow,有两种方式获得:

  • Tensorflow NGC docker
    官方提供的环境中已经预装了 TensorRT/CUDA/cuDNN/cuBLAS/Tensorflow 等必要的软件,并且Tensorflow已经启用了TensorRT, 非常nice,唯一需要注意的是找到和你的环境一致的docker版本就可以。
  • 如果有特殊需求,可以自己构建Tensorflow pip包
    开启TensorRT支持和设置计算能力在前文《Tensorflow c++ so编译》中已有介绍,不再赘述。
    下面是构建tensorflow pip包的命令
bazel build --config=cuda --config=opt //tensorflow/tools/pip_package:build_pip_package
./bazel-bin/tensorflow/tools/pip_package/build_pip_package /tmp/tensorflow_pkg

构建完成后会产生一个类似tensorflow-2.8.0-cp38-cp38-linux_x86_64.whl的问题,pip install即可。
构建过程太漫长,真的不推荐~~~

  • 使用pip install tensorflow 或 tensorflow-gpu 安装的都没有TensorRT支持,F***,看官方介绍是支持的,郁闷了。

2. TF-TRT 转换脚本

import tensorflow as tf
from tensorflow.python.compiler.tensorrt import trt_convert as trt

conversion_params = trt.DEFAULT_TRT_CONVERSION_PARAMS
conversion_params = conversion_params._replace(precision_mode="FP16")
converter = trt.TrtGraphConverterV2(
    input_saved_model_dir=input_saved_model_dir,
    minimum_segment_size=50,
    conversion_params=conversion_params)
converter.convert()
def input_fn():
  Inp1 = np.random.normal(size=(1, 19200, 1080, 3)).astype(np.float32)
  yield (inp1,)
converter.build(input_fn=input_fn)
converter.save(output_saved_model_dir)

注意:

  • 不同Tensorflow版本的TrtGraphConverterV2接口可能是不兼容的,一定要注意你用的TF版本
  • minimum_segment_size:This is the minimum number of nodes required for a subgraph to be replaced by TRTEngineOp能够控制转换出的TRTEngineOp个数,对转换后的TF-TRT模型性能有很大的影响,实际转换过程中需要多次调整,以达到性能最优, 没错,就是在你的部署环境中试出来的。
  • precision_mode 如果精度满足要求当然选低精度,速度不能说翻倍,速度提升70%左右不成问题
  • converter.build 的作用,如果不build, 只进行convert和save的话,会在inference时进行构建 TRTengineOp,一般我们的业务pipline输入数据的shape都是固定的,可以在转换时设置设置输入的shape,构建TRTengineOp,这样可以节省推理时间。
  • 对于 dynamic shape 支持以后再看。

have a good day~~~

Logo

开放原子开发者工作坊旨在鼓励更多人参与开源活动,与志同道合的开发者们相互交流开发经验、分享开发心得、获取前沿技术趋势。工作坊有多种形式的开发者活动,如meetup、训练营等,主打技术交流,干货满满,真诚地邀请各位开发者共同参与!

更多推荐