tensorflow源码解析之common_runtime-device

目录

  1. 核心概念
  2. device
  3. device_factory
  4. device_mgr
  5. device_set

1. 核心概念

在framework部分,我们介绍了DeviceAttributes和DeviceBase两个结构,这些其实是为了我们今天要介绍的Device类做准备的。感兴趣的读者可以去回顾下前面讲过的内容。Device类只是对DeviceBase类的继承,没有添加更多新的数据成员,但提供了Compute计算接口。DeviceSet是一个设备集合类,而DeviceMgr与DeviceSet的不同点在于,它提供了设备管理的功能,为设备查找和计数提供了便利的数据结构。最后,DeviceFactory是为了产生某种类型的设备准备的工厂类,同样的设备类型(比如CPU)会对应不同的工厂,意味着不同的实现,而不同的工厂有着不同的权重。这里的权重是为了辅助我们选择某种类型的设备用的。

2. device

Device类,除了包含对内部私有数据的访问API之外,还包含了核心的计算API Compute,我们先来看一下它的结构:

class Device : public DeviceBase {
  public:
    virtual void Compute(OpKernel* op_kernel, OpKernelContext* context){
        op_kernel->Compute(context);
    }
    virtual void ComputeAsync(AsyncOpKernel* op_kernel, OpKernelContext* context, AsyncOpKernel::DoneCallback done){
        op_kernel->ComputeAsync(context, std::move(done));
    }
    //...
  private:
    const DeviceAttributes device_attributes_;
    DeviceNameUtils::ParsedName parsed_name_;
    OpSegment op_seg_;
    ResourceMgr* rmgr_ = nullptr;
}

TF对于设备名称是有要求的,它必须满足这种格式:/job:_/replica:_/task:_/(gpu|cpu):_,举个例子:/job:train/replica:0/task:3/gpu:2。其中,Device类的数据成员parsed_name_就是对这种设备名称的拆解,感兴趣的读者可以自行看下ParsedName的定义。ResourceMgr和OpSegment我们之前在framework部分也都介绍过了。所以从数据角度讲,Device没有什么新鲜的,只是对原有的关于设备的数据做了一个整合。但从API的角度讲,它包含了一个计算接口Compute,实际上也就是对OpKernel中的Compute接口的封装。

3. device_set

DeviceSet是一个容器类,用于管理一个模型使用的不同设备。这个类相对比较简单,我们看它的结构:

class DeviceSet {
  public:
    //...
  private:
    std::vector<Device*> devices_;
    std::unordered_map<string, Device*> device_by_name_;
    Device* client_device_ = nullptr;
}

其中,device_by_name_是一个从设备全称到设备指针的映射,而client_device_是我们从devices_中挑选的,默认的客户端设备。

4. device_mgr

DeviceMgr顾名思义是一个设备管理类,其实它主要是提供了一系列数据结构来提高API的效率,比如,我们要查找一个给定设备名的设备指针,或者要对某种类型的设备计数。针对这种高频操作,DeviceMgr为其准备了高效的数据结构。类的结构如下:

class DeviceMgr {
  public:
    //...
  private:
    typedef gtl::InlinedVector<Device*, 8> DeviceVec;
    DeviceVec devices_;
    std::unordered_map<StringPiece, Device*, StringPiece::Hasher> device_map_;
    core::Arena name_backing_store_;
    std::unordered_map<string, int> device_type_counts_;
}

device_map_是为了提高查找指定名称的设备的效率,device_type_counts_是为了提高查找指定类型的设备数的效率。

5. DeviceFactory

正如刚才提到过的,DeviceFactory代表了某种设备(比如CPU)的某种实现的工厂类。下面我们看下DeviceFactory类的结构:

class DeviceFactory {
  public:
    static void Register(const string& device_type, DeviceFactory* factory, int priority);
    static DeviceFactory* GetFactory(const string& device_type);
    static Status AddDevices(const SessionOptions& options, const string& name_prefix, std::vector<Device*>* devices);
    static Device* NewDevice(const string& type, const SessionOptions& options, const string& name_prefix);
    virtual Status CreateDevices(const SessionOptions& options, const string& name_prefix, std::vector<Device*>* devices) = 0;
    static int32 DevicePriority(const string& device_type);
};

看完这个类,我们感觉很疑惑,它提供了很多的API,但是没有数据成员,那它注册的那些工厂,存储在哪里呢?

别慌,我们在device_factory.cc文件中,找到了这样的定义:

struct FactoryItem {
    std::unique_ptr<DeviceFactory> factory;
    int priority;
};
std::unordered_map<string, FactoryItem>& device_factories(){
    static std::unordered_map<string, FactoryItem>* factories = new std::unordered_map<string, FactoryItem>;
    return *factories;
}

对于第二个函数,它内部定义了一个静态成员,因此相当于提供了一个全局的从设备类型名称到其生产工厂的映射。每当我们需要这个映射时,就调用这个函数。实际上,DeviceFactory的很多成员函数,就是这样实现的。

另外,TF还提供了一个Registrar类,为DeviceFactory提供了注册的入口:

template<class Factory> class Registrar {
  public:
    explicit Registrar(const string& device_type, int priority=50){
        DeviceFactory::Register(device_type, new Factory(), priority);
    }
};

它实际上是为某种设备类型注册其设备工厂。

关于设备工厂类,我们在代码中经常看到priority,对于权重,我们详细说明一下:

  • 对于同样一种设备类型,不同的注册可以由不同的权重,即同一个设备类型的不同实现,可以拥有不同的权重。权重主要被应用于以下两个方面:
  • (接上)第一,当我们需要为某一个特定的设备类型选择工厂时,拥有最高权重的工厂将会被选择。例如,如果有如下的两种注册信息

    Registrar<CPUFactory1>("CPU", 125);

    Registrar<CPUFactory2>("CPU", 150);

    那么当调用DeviceFactory::GetFactory("CPU")时,CPUFactory2将会被返回。

  • (接上)第二,当需要在DeviceSet中选择一种设备类型时,选择的顺序由权重priority决定。例如,对于以下的两种注册:

    Registrar<CPUFactory>("CPU",100);

    Registrar<GPUFactory>("GPU",200);

    则DeviceType("GPU")将会被优先选择。

  • 不同设备的默认权重如下:GPU:200,SYCL:200,GPUCompatibleCPU:70,ThreadPoolDevice:60,Default:50

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

时间: 2024-10-28 18:08:34

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

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-graph_optimizer

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

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那篇已经见到了.但实际上,问题的难