目录
概述
TensorRT是一个高性能的深度学习推理优化器,可以为深度学习模型提供低延迟、高吞吐率的部署推理,本文旨在从Tensorflow r1.15源码分析它中是如何与TensorRT对接,实现对TF 模型的推理优化
TensorRT Overview
![](https://i-blog.csdnimg.cn/blog_migrate/ba49fda14ea048e6a91f20610c4089e8.jpeg)
TRT C++ API构建模型推理优化与运行的典型时序如上图所示,其中“静态读取模型文件”时序可以参考我之前的文章TensorRT sampleMNIST 详解。而Tensorflow中对接TensorRT则采用的是“动态构建模型”的时序。
架构Overview
![](https://i-blog.csdnimg.cn/blog_migrate/9e104e39ec87aeab950a2f90f6f36712.jpeg)
TRT通过注册优化器到TF框架,从而使得TF在运行过程中能够对计算图进行TRT分割、优化、替换TRT OP节点,最终在执行计算的时候执行TRT OP节点对应的Kernel,从而实现推理优化加速的目的。
![](https://i-blog.csdnimg.cn/blog_migrate/7c3605c9a3f6a617958e9205a4958663.jpeg)
TensorRT在TF框架中的主要功能可以分为上图三步
- 注册TensorRT优化器到TF框架,这是TRT能够在TF框架中起作用的关键入口
- 在TF执行优化器的时候,对计算图进行TensorRT优化
- 在TF执行计算图的时候,调用替换的TRT OP节点执行对应的TRT kernel计算,实现推理优化
TF优化器
TensorFlow的优化器在Grappler模块实现,代码在tensorflow/core/grappler
- tensorflow.grappler.GrapplerItem表示待优化的TensforFlow模型,主要包括计算图、fetch节点、feed节点。
- tensorflow.grappler.Cluster表示可以运行TensorFlow模型的硬件资源集合。一个进程同一时间只能创建一个Cluster.
- tensorflow.grappler.GraphOptimizer是grappler中所有优化类的父类
//tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.cc
...
class VerboseCustomGraphOptimizerRegistrar
: public grappler::CustomGraphOptimizerRegistrar {
public:
VerboseCustomGraphOptimizerRegistrar(
const grappler::CustomGraphOptimizerRegistry::Creator& cr,
const string& name)
: grappler::CustomGraphOptimizerRegistrar(cr, name) {
VLOG(1) << "Constructing a CustomOptimizationPass registration object for "
<< name;
}
};
static VerboseCustomGraphOptimizerRegistrar TRTOptimizationPass_Registrar(
[]() {
VLOG(1)
<< "Instantiating CustomOptimizationPass object TensorRTOptimizer";
return new TRTOptimizationPass("TensorRTOptimizer");
},
("TensorRTOptimizer"));
...
TF框架中通过构建static VerboseCustomGraphOptimizerRegistrar实例注册TRT 优化器TRTOptimizationPass的创建函数到全局map,从而使得在TF执行优化计算图的阶段可以获得TRT优化器对计算图进行TRT图分割、优化、替换
TRT 优化器执行时序