【架构分析】Tensorflow Internals 源码分析5 - TensorRT与Tensorflow

目录

概述

TensorRT Overview

架构Overview

TF优化器

TRT 优化器执行时序

---

概述

TensorRT是一个高性能的深度学习推理优化器，可以为深度学习模型提供低延迟、高吞吐率的部署推理，本文旨在从Tensorflow r1.15源码分析它中是如何与TensorRT对接，实现对TF 模型的推理优化

TensorRT Overview

TRT典型执行时序

TRT C++ API构建模型推理优化与运行的典型时序如上图所示，其中“静态读取模型文件”时序可以参考我之前的文章TensorRT sampleMNIST 详解。而Tensorflow中对接TensorRT则采用的是“动态构建模型”的时序。

架构Overview

TensorRT对接Tensorflow架构图

TRT通过注册优化器到TF框架，从而使得TF在运行过程中能够对计算图进行TRT分割、优化、替换TRT OP节点，最终在执行计算的时候执行TRT OP节点对应的Kernel，从而实现推理优化加速的目的。

TensorRT对接Tensorflow的主流程

 

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 优化器执行时序