Tensorflow、TF-TRT、 TensorRT那些事儿

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

最近一直在折腾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~~~