tensorflow源碼解析之framework-resource


目錄

  1. 什么是resource
  2. 如何使用resource
  3. 如何管理resource
  4. 常用resource
  5. 其它結構
  6. 關系圖
  7. 涉及的文件
  8. 迭代記錄

1. 什么是resource

我們知道,TF的計算是由設備完成的。每個設備包含若干個節點,由這些節點完成實際的計算。有些時候,我們需要在不同的節點之間,共享一些內容,比如,張量值,kv存儲表,隊列,讀取器等等,這些被同設備上的節點共享的內容,就是資源。
所有的資源類都繼承自一個基類,ResourceBase,下面看下它的實現:

class ResourceBase : public core::RefCounted {
public:
    virtual string DebugString() = 0;
    virtual int64 MemoryUsed() const {return 0;};
};

它繼承自core::RefCounted,也就是說,資源本身擁有引用計數的功能。這是為了方便對資源的使用進行監控,在資源無用后能夠及時將它釋放掉。

2. 如何使用resource

普通節點是無法直接使用共享資源的,必須通過一種包含了ResourceOpKernel的特殊節點(關於OpKernel的詳細定義,請參見Kernel,目前僅需要知道,這是節點中用於實際執行計算的類就好了),這種節點幫助普通節點在資源管理器中查找或創建某種特定類型的資源,如圖所示:

graph LR A(OpKernel) -->|派生| B(ResourceOpKernel) B(ResourceOpKernel) -->|創建或查找| C(Resource) A(OpKernel) -->|派生| D(NormalOpKernel) D(NormalOpKernel) -->|使用資源| B(ResourceOpKernel)

ResourceOpKernel類的定義如下(僅保留接口):

template <typename T>
class ResourceOpKernel : public OpKernel {
  public:
    //...
    void Compute(OpKernelContext* context) override LOCKS_EXCLUDED(mu_);
  protected:
    mutex mu_;
    ContainerInfo cinfo_ GUARDED_BY(mu_);//包含了對資源的要求
    T* resource_ GUARDED_BY(mu_) = nullptr;
  private:
    virtual Status CreateResource(T** resource) EXECLUSIVE_LOCKS_REQUIRED(mu_) = 0;//返回一個T派生類的資源對象,這個對象歸當前的ResourceOpKernel對象所有
    virtual Status VerifyResource(T* resource);//校驗resource是否是當前對象需要的類型
    PersistentTensor handle_ GUARDED_BY(mu_);
}

其中的Compute函數,在首次調用時,會根據cinfo_中包含的對資源的要求,從資源管理器中查找,或者創建一個新的資源。
這里牽扯到兩個新的概念,一個是ContainerInfo,一個是資源管理器。我們將在下文中討論。

3. 如何管理resource

最簡單的,可以在設備內維護一個資源名稱到資源實體的映射,但這種管理方式過於粗糙,TF使用容器的概念實現了對資源的分組,下面看一下對容器的定義:

typedef std::pair<uint64,string> Key;
typedef std::unordered_map<Key,ResourceBase*,KeyHash,KeyEqual> Container;

其中,Key是對資源的描述,uint64代表資源類型的哈希值,string代表了資源的名稱,而Container本質上就是一個資源描述到資源指針的映射。
為了對同一個設備上的資源提供統一的管理入口,TF定義了ResourceMgr類,其私有數據成員如下:

class ResourceMgr {
  private:
    //當前設備上的默認容器名稱,如果查找時沒有指定容器,則在默認容器中查找
    const string default_container_;
    mutable mutext mu_;
    //容器名稱到容器指針的映射
    std::unordered_map<string,Container*> containers_ GUARDED_BY(mu_);
    //資源類型的哈希值到資源類型名稱的映射
    std::unordered_map<uint64,string> debug_type_names GUARDED_BY(mu_);
};

資源管理類的公共接口如下:

class ResourceMgr {
public:
    //在container容器中創建一個名為name的資源
    Status Create(const string& container, const string& name, T* resource);
    //在container中查找一個名為name的資源
    Status Lookup(const string& container, const string& name, T** resource) const;
    //如果container中包含名為name的資源,填充到*resource中,否則,使用creater()創建一個資源
    Status LookupOrCreate(const string& container, const string& name, T** resource, std::function<Status(T**)> creater);
    //刪除container中的名為name的資源
    Status Delete(const string& container, const string& name);
    //刪除句柄handle指向的資源
    Status Delete(const ResourceHandle& handle);
    //刪除container中的所有資源,並刪除該container
    Status Cleanup(const string& container);
    //刪除所有容器中的所有資源
    void Clear();
};

注釋很清晰,這里就不再贅述了。
剛才提到,為了讓ResourceOpKernel類的對象能夠找到對應的資源,它在內部包含了一個ContainerInfo的類,它的作用是,幫助一個有狀態的OpKernel決定,它需要通過哪個container/shared_name組合來訪問對應的資源,類的定義如下:

class ContainerInfo {
  public:
    Status Init(ResourceMgr* rmgr, const NodeDef& ndef, bool use_node_name_as_default);
  private:
    ResourceMgr* rmgr_ = nullptr;
    string container_;
    string name_;
    bool resource_is_private_to_kernel_ = false;
};

其中,這個決策過程如下:

  • 如果這個NodeDef的屬性中,包含了"container",就使用這個容器名稱,否則使用rmgr的默認容器;
  • 如果這個NodeDef的屬性中,包含了"shared_name",那么就以此作為資源的名稱,否則,如果"use_node_name_as_default"為真,則使用節點名稱作為資源名稱,如果為假,則生成一個唯一標識作為資源名;

可見,ContainerInfo本質上提供的是,如何幫助節點找到資源的位置,而不是資源位置的靜態描述。為了方便指向資源,TF提出了資源句柄的概念,它的protobuf定義如下:

message ResourceHandleProto {
    //資源所在設備
    string device = 1;
    //資源所在容器
    string container = 2;
    //資源名稱
    string name = 3;
    //資源類型的哈希值
    uint64 hash_code = 4;
    //資源類型名稱,僅調試用
    string maybe_type_name = 5;
}

資源句柄提供了對於資源位置和屬性的詳細描述,通過句柄我們可以唯一定位到一個資源。由於這個資源句柄是可序列化的,因此我們可以方便的傳遞它。另外為了避免使kernel依賴於proto,TF還設計了一個與ResourceHandleProto擁有相同功能的類,ResourceHandle,詳見源代碼。

4. 常用resource

本篇的開頭提到了,常用的資源包括張量值,kv存儲表,隊列,讀取器等等,除了張量值之外,TF為后三者提供了應用的接口,如下:

class LookupInterface : public ResourceBase;
class QueueInterface : public ResourceBase;
class ReaderInterface : public ResourceBase;

其中

  • kv存儲器接口LookupInterface提供了一個批量查找和插入kv存儲器的接口,方便不同的計算節點之間共享kv數據;
  • 隊列接口QueueInterface提供了入隊、出隊等接口,方便不同節點之間共享隊列,這對於計算圖的異步執行是非常重要的,讀取器接口內部也用到了隊列接口;
  • 讀取器接口ReaderInterface是所有TF支持的文件讀取器的統一接口。TF可以支持從不同格式的文件中讀取數據,每一個文件讀取器都是一種資源,都需要繼承自ReaderInterface接口;

其中,關於讀取器接口,其實真正的讀取器資源,並非直接繼承自讀取器接口,而是繼承自它的一個派生類,ReaderBase,這個類為讀取器接口的每一個API提供了一個默認的實現,它的派生類只要重寫這些默認實現,就能完成對不同文件格式的讀取。
同時,為了方便獲取讀取器資源,TF還為ResourceOpKernel生成了一個專用於獲取讀取器接口的派生類ReaderOpKernel:

class ReaderOpKernel : public ResourceOpKernel<ReaderInterface>;

這個類用於查找或者創造讀取器資源,而這些讀取器資源需要派生自ReaderBase類。

5. 其它結構

為了方便對資源進行操作,TF除了設計ResourceOpKernel這個核之外,還設計了兩種核,IsResourceInitialized和ResourceHandleOp,前者用於檢查某個資源是否已經初始化,后者用於為某個資源生成資源句柄。具體定義可以查看源代碼。

6. 關系圖

graph TB A(core::RefCounted)-->|派生|B(ResourceBase) B(ResourceBase)-->|派生|C(LookupInterface) B(ResourceBase)-->|派生|D(QueueInterface) B(ResourceBase)-->|派生|E(ReaderInterface) E(ReaderInterface)-->|派生|F(ReaderBase) F(ReaderBase)-->|派生|G(具體閱讀器資源) H(OpKernel)-->|派生|I(ResourceOpKernel) I(ResourceOpKernel)-->|派生|J(ReaderOpKernel) J(ReaderOpKernel)-.查找或創建.->G(具體閱讀器資源) H(OpKernel)-->|派生|K(IsResourceInitialized) H(OpKernel)-->|派生|L(ResourceHandleOp) L(ResourceHandleOp)-.生成.->M(ResourceHandle) M(ResourceHandle)-.指向.->G(具體閱讀器資源) K(IsResourceInitialized)-.判斷是否初始化.->G(具體閱讀器資源) I(ResourceOpKernel)-.擁有.->N(ContainerInfo) N(ContainerInfo)-.定位.->G(具體閱讀器資源)

7. 涉及的文件

  • resource_handle
  • resource_mgr
  • resource_op_kernel
  • lookup_interface
  • queue_interface
  • reader_interface
  • reader_base
  • reader_op_kernel

8. 迭代記錄

  • v1.0 2018-08-25 文檔創建
  • v2.0 2018-09-08 文檔重構

github地址


免責聲明!

本站轉載的文章為個人學習借鑒使用,本站對版權不負任何法律責任。如果侵犯了您的隱私權益,請聯系本站郵箱yoyou2525@163.com刪除。



 
粵ICP備18138465號   © 2018-2025 CODEPRJ.COM