tensorflow源码解析之common_runtime-graph_optimizer

目录

  1. 核心概念
  2. graph_optimizer
  3. function
  4. optimization_registry

1. 核心概念

本篇主要讲图的优化迭代器。我们在构建原始图的时候,专注于达到目的,但不会去考虑图的执行效率。如果把图的设计过程比喻为高级语言的编写,那么图的优化过程就相当于,将高级语言编译为机器语言的过程中,为了能够加速进行的编译优化。比如,将相同的常数折叠,将Identity节点去除等等。本节主要用来讨论,跟图优化相关的类和函数。

2. graph_optimizer

进行图优化,需要有一个统一的入口,它的输入是图本身,以及图执行的环境,以及优化的配置,输出是优化后的图。这个入口就是GraphOptimizer,我们先来看看它的结构和接口:

class GraphOptimizer {
  public:
    GraphOptimizer(const OptimizerOptions& opts);
    void Optimize(FunctionLibraryRuntime* runtime, Env* env, Device* device, std::unique_ptr<Graph>* graph, const std::unordered_map<const Node*, std::vector<PartialTensorShape>>* shape_map);
  private:
    OptimizerOptions opts_;
};

显然,其中的Optimize就是这个类最重要的API,它将图优化配置opts中的优化过程应用的graph上。可能会将graph替换为另外一个图对象。device是这张图将要运行的设备,它使得优化算法可以考虑针对设备应当考虑的优化选项。shape_map如果是非空的话,它将图中节点的名称映射为部分可知的节点输出形状,可能在某些图优化中会被应用,比如常量折叠优化。

关于图优化,我们需要了解的更为细致一些,所以,先看一下这个类的构造函数具体的实现方式。

GraphOptimizer::GraphOptimizer(const OptimizerOptions& opts) : opts_(opts) {
    if(opts_.opt_level()>=OptimizerOptions::L1){
        opts_.set_do_common_subexpression_elimination(true);
        opts_.set_do_constant_folding(true);
    }
}

通过这个函数我们了解到,优化配置是有级别概念的,当级别大于等于1时,某些默认的优化配置需要被开启,比如“公共子项消除”和“常量折叠”。这些内容我们在具体的优化步骤中也会看到。下面就来看一下核心API,Optimize的内容:

void GraphOptimizer::Optimize(FunctionLibraryRuntime* runtime, Env* env, Device* device, std::unique_ptr<Graph>* graph, const std::unordered_map<const Node*, std::vector<PartialTensorShape>>* shape_map){
    Graph* g = graph->get();
    DumpGraph("Initial",g);//导出当前图的结构

    bool changed = true;
    const int kMaxRounds = 10;
    for(int rounds = 0; rounds < kMaxRounds; ++rounds){
        changed = false;
        if(RemoveListArrayConverter(g)){
            DumpGraph("RemoveListArrayConverter", g);
            changed = true;
        }
        if(opts_.do_function_inlining() && RemoveDeadNodes(g)){
            DumpGraph("RemoveDeadNodes", g);
            changed = true;
        }
        if(opts_.do_function_inlining() && RemoveIdentityNodes(g)){
            DumpGraph("RemoveIdentityNodes", g);
            changed = true;
        }
        if(opts_.do_constant_folding()){
            ConstantFoldingOptions cf_opts;
            cf_opts.shape_map = shape_map;
            bool was_mutated;
            ConstantFold(cf_opts, runtime, env, device, g, &was_mutated).IgnoreError();
            if(was_mutated){
                RemoveDeadNodes(g);
                DumpGraph("ConstFolding",g);
                changed = true;
            }
        }
        if(opts_.do_function_inlining() && FixupSourceAndSinkEdges(g)){
            DumpGraph("FixupSourceAndSinkEdges",g);
            changed = true;
        }
        if(opts_.do_common_subexpression_elimination() && OptimizeCSE(g,nullptr)){
            DumpGraph("ExpandInlineFunctions",g);
            changed = true;
        }
        if(!changed) break;
    }

    //由于flib_def永远不会消失,因此我们可以放心的使用它来构建新图
    std::unique_ptr<Graph> copy(new Graph(g->flib_def()));
    CopyGraph(*g, copy.get());
    graph->swap(copy);

    DumpGraph("ReCopy", graph->get());
}

在对图进行优化时,我们不可能一蹴而就的,因为优化之间会相互影响,比如我们对图进行了A优化,对于A优化来说,此时图已经是最优的了,但之后我们又对图进行了B优化,此时对于B优化来说,图已经是最优的了,但对于A优化来说则未必。因此图优化是一个循环上升的过程,TF设置了最高的优化是10遍,对于大多数图来说,也就足够了。

在图优化的过程中,我们发现了很多之前没见过的函数,这些函数的定义都在function.h文件中,为了加深对于图优化过程的理解,下面我们了解下这个文件中的函数。

3. function

function.h文件中,没有类定义,全部都是硬生生的函数定义,干货满满。

//kernel生成器,根据FunctionLibraryRuntime和NodeDef来生成kernel
typedef std::function<Status(FunctionLibraryRuntime*, const NodeDef&, std::unique_ptr<OpKernel>*)> CustomKernelCreator;
void RegisterDefaultCustomKernelCreator(CusteomKernelCreator cb);//kernel生成器的注册器

//创建一个FunctionLibraryRuntime,用来实例化lib_def中的函数,并在device上运行,如果custom_kernel_creator是非空的,它会被返回的runtime用来生成kernel
std::unique_ptr<FunctionLibraryRuntime> NewFunctionLibraryRuntime(const DeviceMgr* device_mgr, Env* env, Device* device, int graph_def_version, const FunctionLibraryDefinition* lib_def, const OptimizerOptions& optimizer_options, CusteomKernelCreator custom_kernel_creator);

//与之前的函数类似,只不过返回的runtime直接利用RegisterDefaultCustomKernelCreator注册的全局custom_kernel_creator来生成新的kernel
std::unique_ptr<FunctionLibraryRuntime> NewFunctionLibraryRuntime(const DeviceMgr* device_mgr, Env* env, Device* device, int graph_def_version, const FunctionLibraryDefinition* lib_def, const OptimizerOptions& optimizer_options);

//函数体的内容
struct FunctionBody {
    FunctionDef fdef;
    Graph* graph = nullptr;
    DataTypeVector arg_types;
    DataTypeVector ret_types;
    gtl::InlinedVector<Node*, 4> arg_nodes;
    gtl::InlinedVector<Node*, 4> ret_nodes;

    FuntionBody(){}
    FunctionBody(const FunctionDef& f, DataTypeSlice arg_types, DataTypeSlice ret_types, Graph* g);
    ~FunctionBody();
};

//删除以下节点,第一,无状态的,第二,无参数的,第三,对输出无贡献的
bool RemoveDeadNodes(Graph* g);

//寻找如下的模式,src-(in)->node-(out)->dst,如果node是identity节点,in是唯一的输入数据边,out是唯一的输出数据边,则使用src->dst重写以上模式
bool RemoveIdentityNodes(Graph* g);

//将图中的_ListToArray和_ArrayToList转化为Identity节点
bool RemoveListArrayConverter(Graph* g);

//对于图中的每个节点,如果lib指明这个节点是一个函数调用,那么内联这个函数体。如果至少一个节点被内联了,返回true。
bool ExpandInlineFunctions(FunctionLibraryRuntime* lib, Graph* graph);

//将graph中的内容导出到日志文件,如果日志级别足够高的话
void DumpGraph(StringPiece label, const Graph* g);

//应用图重写的优化,例如内联、死节点移除等
void OptimizeGraph(FunctionLibraryRuntime* lib, std::unique_ptr<Graph>* g);

//将一个函数的图转化为GraphDef
void ToGraphDef(const Graph* g, GraphDef* gdef, bool pretty = false);

//给定一个数值函数,返回它的导数函数
FunctionBody* SymbolicGradient(const FunctionBody& f);

//将一个FunctionDef示例化为一个graph,设置fbody指向拥有FunctionDef的FunctionBody
Status FunctionDefToBodyHelper(const FunctionDef& fdef, const AttrSlice& attrs, const FunctionLibraryDefinition* const lib_def, const std::function<Status(const string&, const OpDef**)>& get_func_sig, FunctionBody** fbody);

现在回过头来看GraphOptimizer类中的Optimize函数,首先它把Array和List相互转换节点变为Identity节点,然后删除了死节点,删除Identity节点,进行常量折叠,修复输入输出边,进行公共子项消除,最终完成了对图的优化。

4. optimization_registry

optimization_registry.h文件中,包含了一些维护一个全局的图优化遍历注册器所需要的类,在会话初始化一张图时,会使用这个全局优化遍历注册器来对图进行优化。

首先我们来看第一个类,GraphOptimizationPassOptions,顾名思义,它包含了图优化遍历所需要的参数。这些足够作为一个字典的键值,我们通常会使用一个字典来保持各个图优化遍历器的状态。

struct GraphOptimizationPassOptions {
    string session_handle;
    const SessionOptions* session_options = nullptr;
    const CostModel* cost_model = nullptr;
    FunctionLibraryDefinition* flib_def = nullptr;
    const DeviceSet* device_set = nullptr;
    //如果优化遍历在图分割之前被使用,那么它优化的对象就是这个graph,如果是图分割之后被使用,那么这个graph是null
    std::unique_ptr<Graph>* graph = nullptr;
    //进行图分割后的优化遍历时使用
    std::unordered_map<string, std::unique_ptr<Graph>* partition_graphs = nullptr;
};

图优化遍历,按照在图分割之前还是之后进行,可以分为两类,但我们使用了GraphOptimizationPassOptions这样一个接口。

接下来是GraphOptimizationPass类,所有的图优化遍历类,都是这个类的子类,它的结构也非常简单。

class GraphOptimizationPass {
  public:
    virtual ~GraphOptimizationPass() {}
    virtual Status Run(const GraphOptimizationPassOption& options) = 0;
};

当我们拥有了多种图优化遍历的算法之后,需要对这些进行统一管理,因此TF提出了一种对图优化遍历算法进行统一注册和管理的类:

//这里的键值为phase,图优化遍历算法是按照phase的升序顺序执行的,在一个phase内部,执行顺序是未定义的
typedef std::map<int, std::vector<std::unique_ptr<GraphOptimizationPass>>> GraphOptimizationPasses;

class OptimizationPassRegistry {
  public:
    enum Grouping {
        PRE_PLACEMENT,//在cost model赋值之后,在节点放置算法之前
        POST_PLACEMENT,//在节点放置算法之后
        POST_REWRITE_FOR_EXEC,//在利用feed/fetch节点进行重写之后
        POST_PARTITIONING,//在图分割之后
    };
    void Register(Grouping grouping, int phase, std::unique_ptr<GraphOptimizationPass> pass);//注册图优化遍历算法
    Status RunGrouping(Grouping grouping, const GraphOptimizationPassOptions& options);//运行一个groupping中所有的图优化遍历算法,按照phase的升序运行
    static OptimizationPassRegistry* Global();//返回一个全局的图优化遍历注册器
  private:
    std::map<Grouping, GraphOptimizationPasses> groups_;
};

总结一下,groups是一个双层的映射,先从Grouping映射到图优化遍历算法组,这个算法组本身也是个映射,从phase映射到真正的图优化遍历算法,如下:

graph LR
Grouping-->GraphOptimizationPasses
phase-->GraphOptimizationPass

最后,TF为刚才的注册器提供了一个全局的入口:

class OptimizationPassRegistration {
  public:
    OptimizationPassRegistration(OptimizationPassRegistry::Grouping grouping, int phase, std::unique_ptr<GraphOptimizationPass> pass){
        OptimizationPassRegistry::Global->Register(grouping,phase,std::move(pass));
    }
};

原文地址:https://www.cnblogs.com/jicanghai/p/9569938.html

时间: 2024-09-30 01:22:26

tensorflow源码解析之common_runtime-graph_optimizer的相关文章

tensorflow源码解析系列文章索引

文章索引 framework解析 resource allocator tensor op node kernel graph device function shape_inference common_runtime解析 device session graph_optimizer executor-1 executor-2 direct_session 后记 关于起源 阅读tensorflow源码时,为了敦促自己主动思考,把阅读的笔记整理成了博客,拿出来跟大家分享. 关于迭代 文章都是工作

Tensorflow源码解析1 -- 内核架构和源码结构

1 主流深度学习框架对比 当今的软件开发基本都是分层化和模块化的,应用层开发会基于框架层.比如开发Linux Driver会基于Linux kernel,开发Android app会基于Android Framework.深度学习也不例外,框架层为上层模型开发提供了强大的多语言接口.稳定的运行时.高效的算子,以及完备的通信层和设备层管理层.因此,各大公司早早的就开始了深度学习框架的研发,以便能占领市场.当前的框架有数十种之多,主流的如下(截止到2018年11月) 显然TensorFlow是独一无

tensorflow源码解析之common_runtime-executor-上

目录 核心概念 executor.h Executor NewLocalExecutor ExecutorBarrier executor.cc structs GraphView ExecutorImpl ExecutorState details 1. 核心概念 执行器是TF的核心中的核心了,前面做了这么多的准备工作,最后要在这里集大成了,想想还有点小激动.不过笔者在这里先打个预防针,执行器的概念多.结构复杂,想要透彻理解并不容易,为了保持文章的易读性,我们也是尽量对细枝末节做了舍弃,以求反

tensorflow源码解析之framework-allocator

目录 core/framework resource allocator 核心概念 给出的只是内存分配器的接口,没有给出具体实现. Allocator Allocator是一个内存分配器的接口类,它规定了一个内存分配器需要具有哪些API.具体看代码: class Allocator { public: virtual void* AllocateRaw(size_t alignment, size_t num_bytes) = 0; virtual void DeallocateRaw(void

tensorflow源码解析之framework-node

目录 核心概念 node_def 1. 核心概念 TF中的图由节点构成,每个节点包含了一个操作,表名这个节点的作用,比如,将两个输入矩阵相乘,输出结果.节点是自带图结构的,每个节点都包含了输入的来源,因此若干节点的集合就能无需其它信息的生成一张图.节点必须被放置在某一个设备上,为了减少跨设备传输数据,也为了提高计算效率,TF还专门开发了相应的节点放置算法. 2. node_def 我们先来看一下节点的定义: message NodeDef { string name = 1;//节点名称 str

tensorflow源码解析之framework-function

目录 核心概念 FunctionDef function related classes 1. 核心概念 在讲解function的概念之前,我们要先回顾下op.op是规定了输入和输出的操作声明,在研究node的时候我们也看到,NodeDef是包含OpDef的,那么是不是op就只能是节点级别的操作呢?并非如此,操作是可以嵌套的,也就是说,操作A可能内部包含了操作BCD.从这个角度理解function就容易了,function其实就是一些大的op.函数的本质是给定输入,经过计算给出输出,这与op的定

tensorflow源码解析之common_runtime-device

目录 核心概念 device device_factory device_mgr device_set 1. 核心概念 在framework部分,我们介绍了DeviceAttributes和DeviceBase两个结构,这些其实是为了我们今天要介绍的Device类做准备的.感兴趣的读者可以去回顾下前面讲过的内容.Device类只是对DeviceBase类的继承,没有添加更多新的数据成员,但提供了Compute计算接口.DeviceSet是一个设备集合类,而DeviceMgr与DeviceSet的

tensorflow源码解析之common_runtime-executor-下

目录 核心概念 executor.h Executor NewLocalExecutor ExecutorBarrier executor.cc structs GraphView ExecutorImpl ExecutorState details 3.4 ExecutorState 在执行器的执行图计算的时候,需要一个结构来保存当前计算的即时信息,TF为此设计了类ExecutorState,它被用来保存每一个对ExecutorImpl::Run调用的状态信息.它会在一个节点已经准备好之后调度

tensorflow源码解析之common_runtime-direct_session

目录 核心概念 direct_session direct_session.h direct_session.cc 1. 核心概念 读过之前文章的读者应该还记得,session是一个执行代理.我们把计算图和输入交给session,由它来调度执行器,执行计算产生结果.TF给我们提供了一个最简单的执行器direction_session.按照当前的理解,我们觉得direction_session的实现应该是非常简单而直接的,毕竟执行器的复杂结构我们在executor那篇已经见到了.但实际上,问题的难