Tensorflow、TF-TRT、 TensorRT那些事儿

小白龙呢

已于 2022-05-16 10:43:26 修改

阅读量3.8k

点赞数 2

文章标签： tensorflow 深度学习机器学习

于 2022-05-13 15:40:20 首次发布

本文链接：https://blog.csdn.net/wohenibdxt/article/details/124748731

版权

cuda 专栏收录该内容

4 篇文章

订阅专栏

本文介绍了如何利用TensorFlow的TF-TRT接口将训练好的模型进行优化，以提高GPU上的推理性能。重点讨论了TensorFlow、TensorRT的基本概念，以及TF-TRT的转换过程，包括启用TensorRT的TensorFlow环境搭建、模型转换脚本的关键参数设置。通过转换和优化，可以实现低延迟和高吞吐量的推理。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

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

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