tensorflow源碼解析之framework拾遺

阿新 • • 發佈：2018-09-13

節點 repeat 普通相互 eve pen lib ash 除了

把framework中剩余的內容，按照文件名進行了簡單解析。時間原因寫的很倉促，算是占個坑，後面有了新的理解再來補充。

allocation_description.proto

一個對單次內存分配結果進行信息描述的proto。

attr_value

之前在講op的時候提到過，操作是有參數的。而AttrValue表示的就是參數的值。先看一下它的proto定義：

message AttrValue {
    message ListValue {
        repeated bytes s = 2;
        repeated int64 i = 3;
        repeated float f = 4;
        repeated bool b = 5;
        repeated DataType type = 6;
        repeated TensorShapeProto shape = 7;
        repeated TensorProto tensor = 8;
        repeated NameAttrList func = 9;
    }
    
    oneof value {
        bytes s = 2;
        int64 i = 3;
        float f = 4;
        bool b = 5;
        DataType type = 6;
        TensorShapeProto shape = 7;
        TensorProto tensor = 8;
        ListValue list = 1;
        
        //func代表一個函數，func.name代表一個函數的名稱，或者核心操作的名稱，func.attr.first是為函數定義的參數的名稱，func.attr.second是上述參數的值
        NameAttrList func = 10;
        
        //placeholder僅在函數內部的節點中使用，它意味著，這個屬性值直到函數被初始化時才會提供。例如，我們假設函數FN中有一個節點N，節點N擁有屬性A，A的屬性值是"foo"，如果FN在初始化時將foo設置為bar，那麽N節點的A屬性的值也已被設置為bar
        string placeholder = 9;
    }
}
message NameAttrList {
    string name = 1;
    map<string, AttrValue> attr = 2;
}

可見，操作屬性的可取值類型很豐富，可以是字符串、整型、浮點型、布爾型、元類型（指代一種數據類型）、張量形狀、張量、列表、函數、填充值（placeholder）等等。

attr_value_util中包含了一些方便對參數值進行設置和表示的輔助函數。

bfloat16

Google發現用16位精度的浮點數，代替32位精度的浮點數，進行神經網絡的計算，模型的精度並沒有明顯下降，但模型大小明顯降低，因此TF推出了這種16位的浮點數。

熟悉單精度浮點數表示的朋友應該知道，單精度float的32位表示是由1位符號位+8位階碼+23位尾數組成的。IEEE也提出過一個16位精度的浮點數，但與目前的32位精度浮點數之間的轉換比較復雜，因此TF設計了一個新的16位浮點數表示，1位符號位+8位階碼+7位尾數。跟float的符號位和階碼是相同的，僅尾數位數不同，因此方便了相互轉換。

cancellation

在計算圖的執行過程中，如果需要臨時終止計算圖的執行，比如發現輸入填寫錯誤，或者編碼錯誤，計算圖並不會馬上停下來，因為一方面有很多異步運算正在進行，另一方面很多計算是在遠程設備上執行的，我們必須通知到正在執行的遠程設備。這些操作就需要一個實體來負責，這就是CancellationManager。

class CancellationManager {
  public:
    //...
    void StartCancel();//運行所有跟當前的取消管理器有關的回調函數
    bool IsCancelled();//當且僅當StartCancel被調用後，返回true
    CancellationToken get_cancellation_token();//在註冊和解註冊回調函數時，需要用到的令牌
    bool RegisterCallback(CancellationToken token, CancellCallback callback);//在一個token上註冊取消回調函數
    bool DeregisterCallback(CancellationToken token);//在一個token上解註冊取消回調函數
  private:
    bool is_cancelling_;
    std::atomic_bool is_cancelled_;
    
    mutex mu_;
    Notification cancelled_notification_;
    CancellationToken next_cancellation_token_ GUARDED_BY(mu_);
    gtl::FlatMap<CancellationToken, CancelCallback> callbacks_ GUARDED_BY(mu_);
};

common_shape_fns

包含了一些通用的，形狀推斷函數中可能會用到的功能函數，比如，卷積運算，怎樣通過輸入形狀，核大小，padding大小，stride大小來判斷輸出的形狀。

control_flow

控制流是一個大問題，TF官方有一篇文章解釋的很好，後續會專門為此寫一篇博文。當前需要理解的是，在TF中為了實現控制流（條件、循環等），需要給張量添加一些附加的屬性。比如，考慮如下的代碼：

int fun(int t){
    int s = 0;
    for(int i=1;i<=100;++i){
        s += t;
    }
    return s;
}

int x(){
    return fun(2);
}

int y(){
    return fun(3);
}

這段代碼如果放到TF計算圖中實現，同樣的一個變量s，會有多個不同的值，在fun函數內部，不同叠代輪次中s的值是不同的，在x函數中調用，和在y函數中調用，相同輪次s的值也是不同的。為了對不同調用，不同叠代輪次內的s做區分，TF提出了一個結構體：

struct FrameAndIter {
    uint64 frame_id = kIllegalFrameId;
    int64 iter_id = kIllegalIterId;
    
    FrameAndIter(){}
    FrameAndIter(uint64 frame, int64 iter){
        frame_id = frame;
        iter_id = iter;
    }
    
    bool operator==(const FrameAndIter& other) const {
        return (frame_id == other.frame_id && iter_id == other.iter_id);
    }
};

仔細看過之前博文的朋友已經看出來，這個結構跟執行器中的TaggedNode很像，只不過FrameAndIter針對張量，TaggedNode針對節點，感興趣的讀者可以去回顧下executor-下。

cost_graph

在對計算圖進行優化時，一個很重要的信息，是對計算圖中各節點的計算消耗進行估計。TF專門提出了一個統計計算圖消耗（內存消耗、計算時間消耗）的模型，下面看一下它的結構：

message CostGraphDef {
    message Node {
        string name = 1;
        string device = 2;
        int32 id = 3;
        
        message InputInfo {
            int32 preceding_node = 1;
            int32 preceding_port = 2;
        }
        repeated InputInfo input_info = 4;
        
        message OutputInfo {
            int64 size = 1;
            int64 alias_input_port = 2;
            TensorShapeProto shape = 3;
            DataType dtype = 4;
        }
        repeated OutputInfo output_info = 5;
        
        int64 temporary_memory_size = 6;//臨時內存損耗
        int64 host_temp_memory_size = 10;//host臨時內存損耗
        int64 device_temp_memory_size = 11;//device臨時內存損耗
        int64 host_persistent_memory_size = 12;//host永久內存損耗
        int64 device_persistent_memory_size = 16;//device永久內存損耗
        
        int64 compute_cost = 9;//該節點計算時長的估計，單位毫秒
        int64 compute_time = 14;//純計算損耗，不包括內存訪問損耗
        int64 memory_time = 15;//內存訪問損耗，不包含計算損耗
        bool is_final = 7;//當前節點的輸出，是否是整個計算圖的輸出，如果是，則這個輸出不能被拋棄
        
        repeated int32 control_input = 8;//當前節點的控制輸入
    }
    repeated Node node = 1;
}

可見，對計算圖損耗的統計，就是對節點損耗統計的集合。而對於節點，除了基礎的屬性信息和輸入輸出信息之外，主要包含了內存消耗和時間消耗的信息。

fake_input

為了測試NodeDefBuiler的功能，我們需要為節點準備一些輸入，但其實大部分功能的測試並不需要真實的數據，我們只需要有一個輸入的形式在。因此TF推出了FakeInput結構，它是一個內部使用的結構，具體實現是FakeInputImpl，感興趣的讀者可以去看下源碼。

load_library

當運行時環境初始化的時候，需要把op和kernel的定義載入內存。首次載入時，這些資源會被放入一個全局的數據結構中，後續需要時可以從中檢索。

Status LoadLibrary(const char* library_filename, void** result, const void** buf, size_t* len){
    static mutex mu;
    static std::unordered_map<string, Library> loaded_libs;
    //...
}

也就是說，被載入的庫其實被存在一個全局的map中，其中key為庫所在的文件名，value為一個Library結構，下面看下它的定義：

struct Library {
    void* handle = nullptr;
    OpList op_list;
};

即，這裏的庫實際上包含的是一個指向自身的句柄，以及一個操作的集合。

log_memory

在程序運行時，我們經常需要分配內存空間，內存的分配通常被分為兩種情況，第一是在OpKernel計算時分配內存，這些分配會由一個進程級別的編號（step_id）來標識，第二是各種特殊的內存分配場景，包括：

當進行即時的常量折疊優化時；
當進行OpKernel的構建時；
當使用外部代碼，比如C API分配張量時；
當為網絡傳輸分配內存時；
當為GPU傳輸過來的proto分配內存時；
當調用者並沒有指明step_id時；

了解了哪些情況需要記錄內存分配，還需要知道，在不同的情況下需要記錄哪些信息，為了區分內存分配信息記錄的不同場景，TF做出了如下分類：

記錄普通張量的內存分配，普通張量包括OpKernel計算時申請的內存，以及上述特殊情況；
記錄普通張量的內存回收；
當把某個張量作為輸出時，需要記錄；
原始內存的分配，包括的場景有Eigen內存的分配，內存拷貝；
原始內存的回收；

有了這些基本概念，理解內存分配的相關結構就容易了，首先我們看下，TF為5種內存分配記錄的情況，設計的proto：

//在哪一步進行了內存分配
message MemoryLogStep {
    int64 step_id = 1;//進程級別的步驟id，進程內部相同，進程之間不同
    string handle = 2;//描述當前步輸入和輸出的句柄
};

//張量內存分配的信息
message MemoryLogTensorAllocation {
    int64 step_id = 1;
    
    //進行內存分配的kernel名稱，比如"/affine2/weights/Assign"
    string kernel_name = 2;
    TensorDescription tensor = 3;//分配的張量的細節
};

//張量內存回收的信息
message MemoryLogTensorDeallocation {
    int64 allocation_id = 1;
    string allocator_name = 2;
};

//張量設置為輸出的信息
message MemoryLogTensorOutput {
    int64 step_id = 1;
    string kernel_name = 2;
    int32 index = 3;//被設置的輸出的索引
    TensorDescription tensor = 4;
};

//原始的內存分配信息
message MemoryLogRawAllocation {
    int64 step_id = 1;
    string operation = 2;
    int64 num_bytes = 3;
    uint64 ptr = 4;
    int64 allocation_id = 5;
    string allocator_name = 6;
};

//原始的內存回收信息
message MemoryLogRawDeallocation {
    int64 step_id = 1;
    string operation = 2;
    int64 allocation_id = 3;
    string allocator_name = 4;
    bool deferred = 5;
};

其次，我們看下內存分配的功能類LogMemory：

class LogMemory {
  public:
    static bool IsEnabled();
    static void RecordStep(int64 step_id, const string& handle);
    static void RecordTensorAllocation(const string& kernel_name, int64 step_id, const Tensor& tensor);
    static void RecordTensorDeallocation(int64 allocation_id, const string& allocator_name);
    static void RecordTensorOutput(const string& kernel_name, int64 step_id, int index, const Tensor& tensor);
    static void RecordRawAllocation(const string& operation, int64 step_id, size_t num_bytes, void* ptr, Allocator* allocator);
    static void RecordRawDeallocation(const string& operation, int64 step_id, void* ptr, Allocator* allocator, bool deferred);
};

可以看到，主要的API基本上跟前面的proto一一對應。
這些API內部如何實現的呢？我們挑一個最簡單的來看下：

void LogMemory::RecordStep(const int64 step_id, const string& handle){
    MemoryLogStep step;
    step.set_step_id(step_id);
    step.set_handle(handle);
    OutputToLog(step);
}

可見，這些API的主要作用就是把輸入參數放入相應的proto，然後以日誌的形式將這些proto輸出，具體的函數如下：

template <typename T>
void OutputToLog(const T& proto){
    string type_name = proto.GetTypeName();
    const size_t index = type_name.find_last_of(".");
    if (index != string::npos) type_name = type_name.substr(index + 1);
    LOG(INFO) << LogMemory::kLogMemoryLabel << " " << type_name << " { " << ProtoShortDebugString(proto) << " }";
}

memory_types

提供了一個，根據NodeDef，獲取節點輸入輸出內存類型的函數，接口如下：

Status MemoryTypesForNode(const OpRegistryInterface* op_registry, const DeviceType& device_type, const NodeDef& ndef, MemoryTypeVector* input_memory_types, MemoryTypeVector* output_memory_types);

numeric_op

定義了四類最常見的數值操作類型：

單輸入單輸出，輸入輸出的類型相同，比如自增運算；
雙輸入單輸出，相同類型，比如標量加法；
輸入和輸出擁有相同的形狀，且輸入輸出一一對應，比如矩陣元素翻倍運算；
輸入和輸出擁有相同的形狀，且兩個輸入對應一個輸出，比如兩個矩陣相加運算；

代碼中，這四種運算的定義如下：

class UnaryOp : public OpKernel;
class BinaryOp : public OpKernel;
template <class T, class CHILD> class UnaryElementWiseOp : public UnaryOp<T>;
template <class T, class CHILD> class BinaryElementWiseOp : public BinaryOp<T>;

numeric_types

定義了常用的數值類型。

register_types

在定義一個操作的時候，我們往往會提供一個類型參數，但一方面不一定所有的操作都支持所有的數據類型，另一方面，當前的硬件也不一定支持所有的數據類型。因此有必要為操作設計一個可以快速實例化為各數據類型具體操作的宏，也有必要為不同的硬件設計不同的可用宏。

因此，這裏的宏分為兩類，一類是TF_CALL_float這種針對具體數據類型的具體宏，另一類是TF_CALL_ALL_TYPES這種類組宏，第二類宏通過調用第一類宏來工作。例如：

#define TF_CALL_INTEGRAL_TYPES(m)   TF_CALL_int64(m) TF_CALL_int32(m) TF_CALL_uint16(m) TF_CALL_int16(m) TF_CALL_uint8(m) TF_CALL_int8(m)

register_types_traits

這個文件的功能是，在由CPU向GPU拷貝數據時，為POD數據類型提供代理類型。什麽POD數據類型呢？POD的全程是Plain Old Data，簡單來說，一個類或者結構體，在經過二進制拷貝後還能保持數據不變，這就是POD數據類型。在由CPU向GPU拷貝數據時，實際上拷貝的是二進制的數據，因此它們實際的數據類型，在傳輸時可以忽略掉，直接拷貝二進制數據就好了。

rendezvous

一個Rendezvous是一個從生產者向消費者傳輸張量的抽象，它由一個通道映射表組成。每一個通道由一個Rendezvous鍵唯一標識，這個鍵由"producer,consumer"組成，生產者和消費者都是TF中的設備。

生產者通過調用Send()函數，將一個張量通過通道傳遞給消費者，消費者通過調用Recv()函數，從一個通道中接收傳遞過來的張量。消費者按照生產者生產的順序接收傳輸的張量。

消費者可以在生產者將張量生產出來之前或之後，索要這個張量。消費者可以選擇進行一個阻塞調用，或者提供一個回調函數，在任何一種情況下，只要張量生產出來，消費者都會第一時間得到它。生產者從不阻塞。

由於比較簡單，我們僅列出API的名稱，具體簽名和實現，大家可以參考源代碼：

class Rendezvous : public core::RefCounted {
  public:
    static string CreateKey(...);//創建一個傳輸的鍵
    static Status ParseKey(...);//解析一個傳輸的鍵
    virtual Status Send(...) = 0;
    virtual void RecvAsync(...) = 0;
    Status Recv(...);
    virtual void StartAbort(...) = 0;
}

session_state

這裏面包含了兩個類，SessionState保存了我們需要在不同運行中保存的張量，比如深度學習需要叠代的訓練，每次訓練之間需要共享一些數據，就保存在這裏。而TensorStore保存了我們在當前運行中需要共享的張量，它被所有的op_kernel共享。

class SessionState {
  public:
    Status GetTensor(const string& handle, Tensor* tensor);
    Status AddTensor(const string& handle, const Tensor& tensor);
    Status DeleteTensor(const string& handle);
    int64 GetNewId();
    static const char* kTensorHandleResourceTypeName;
  private:
    mutex state_lock_;
    int64 tensor_id_ = 0;
    std::unordered_map<string, Tensor> tensors_;
};
class TensorStore {
  public:
    struct TensorAndKey {
        Tensor tensor;
        int64 id;
        string device_name;
        string GetHandle(const string& tensor_name){
            return strings::StrCat(tensor_name, ";", id, ";", device_name);
        }
    };
    Status AddTensor(const string& name, const TensorAndKey& tk);
    Status SaveTensors(const std::vector<string>& output_names, SessionState* session_state);
  private:
    mutex lock_;
    std::unordered_map<string, TensorAndKey> tensors_ GUARDED_BY(lock_);
};

step_stats

主要是關於運行時間和內存使用的統計。單步統計由設備統計構成，而設備統計由節點統計構成。下面看下它們的核心結構：

message NodeExecStats {
    string node_name = 1;
    int64 all_start_micros = 2;
    int64 op_start_rel_micros = 3;
    int64 op_end_rel_micros = 4;
    int64 all_end_rel_micros = 5;
    repeated AllocatorMemoryUsed memory = 6;
    repeated NodeOutput output = 7;
    string timeline_label = 8;
    int64 scheduled_micros = 9;
    uint32 thread_id = 10;
    repeated AllocationDescription referenced_tensor = 11;
    MemoryStats memory_stats = 12;
};

message DeviceStepStats {
    string device = 1;
    repeated NodeExecStats node_stats = 2;
};

message StepStats {
    repeated DeviceStepStats dev_stats = 1;
};

剩余的AllocatorMemoryUsed，NodeOutput，AllocationDescription，MemoryStats，比較簡單，大家感興趣可以直接去看源碼。

summary.proto

用在Tensorboard中，用於顯示每個節點元素的匯總信息。

type_index

c++中的std::type_index提供了RTTI，即運行時類型識別的功能。它本質上包含了一個類型的哈希值，可以作為一個無序映射的鍵。cppreference上的這段示例代碼非常清楚：

#include <iostream>
#include <typeinfo>
#include <typeindex>
#include <unordered_map>
#include <string>
#include <memory>
 
struct A {
    virtual ~A() {}
};
 
struct B : A {};
struct C : A {};
 
int main()
{
    std::unordered_map<std::type_index, std::string> type_names;
 
    type_names[std::type_index(typeid(int))] = "int";
    type_names[std::type_index(typeid(double))] = "double";
    type_names[std::type_index(typeid(A))] = "A";
    type_names[std::type_index(typeid(B))] = "B";
    type_names[std::type_index(typeid(C))] = "C";
 
    int i;
    double d;
    A a;
 
    // 註意我們正在存儲指向類型 A 的指針
    std::unique_ptr<A> b(new B);
    std::unique_ptr<A> c(new C);
 
    std::cout << "i is " << type_names[std::type_index(typeid(i))] << ‘\n‘;
    std::cout << "d is " << type_names[std::type_index(typeid(d))] << ‘\n‘;
    std::cout << "a is " << type_names[std::type_index(typeid(a))] << ‘\n‘;
    std::cout << "b is " << type_names[std::type_index(typeid(*b))] << ‘\n‘;
    std::cout << "c is " << type_names[std::type_index(typeid(*c))] << ‘\n‘;
}

這裏，std::type_index(typeid(i))實際上是通過std::type_index的構造函數，構造了一個TypeIndex對象。雖然b和c被定義為指向A的指針，但通過typeid我們仍然能辨識出它實際上指向的是什麽類型。

這種附加的類型信息是有代價的。在某些平臺上，我們希望通過避免掉這種類型信息，來獲得更小的二進制存儲空間。因此TF提出了一個簡化版的TypeIndex類，它模擬了std::type_index的功能，但是沒有使用RTTI信息，因此也就不能提供真正的類型信息，只是返回一個標誌，表示RTTI已被禁用。類中包含的哈希碼對每個類是唯一的，然而它是在運行時產生的，因此這個哈希值被序列化後並沒有意義，因為每次運行的時候這個哈希值都不一定相同。

下面我們來看下TF自定義的TypeIndex的實現：

class TypeIndex {
  public:
    TypeIndex(const TypeIndex& src) : hash_(src.hash_){}
    //...
  private:
    TypeIndex(const uint64 hash) : hash_(hash){}
    uint64 hash_;
};

types

關於內存類型MemoryType和設備類型DeviceType的定義，還有一些瑣碎的信息，詳見源碼。

type_traits

定義了一些常用類型判斷的模板，主要包含以下四種：

is_quantized，是否是quantized type；
is_complex，是否是復數類型；
is_simple_type，是否是簡單類型；
is_signed，是否是帶符號的類型；

unique_tensor_references

一組唯一的tensor references的集合。在向這個集合中加入tensor的時候會判斷，指向這個tensor內部的buffer的引用是否已經存在的，如果已存在，不做任何操作，如果不存在，則插入。

這裏做了一個小小的優化，因為這個類中存儲的不同的張量引用不會太多，所以一開始可以用一個最大長度為4的內聯向量存儲這些引用。當不同的張量引用數超過4時，使用一個set來存儲。這樣在大部分情況下，都能保證較好的插入和查詢性能。這種優化方式比較普遍，在TF中也經常被使用。

variant

一個類型擦除的容器，可用於存儲各種數據類型。實現方式與std::any很像，但對於存儲的數據類型有限制，只能存儲有限幾個類型的數據。它能存儲的數據必須滿足以下的幾個條件：

這個類是可以復制構造的；
類有默認構造函數；
它要麽是一個protobuf，要麽是一個TF的tensor，要麽定義了以下的三種函數

    string TypeName() const;
    void Encode(VariantTensorData* data) const;
    void Decode(const VariantTensorData& data);

使用get

Variant對象將序列化和反序列化的工作交給內部的數據類型去做，對於一些未實現序列化和反序列化功能的POD類型（plain old data），TF的tensor類型，以及protobuf類型，TF在文件中提供了一些輔助函數。如果需要在Variant中存放其它數據類型，需要單獨提供Encode和Decode實現。

在Variant中存儲的數據類型，通常會包含指向其它TF的tensor的引用。為了有效的支持這種情況，TF給序列化的結構提供了明確的結構，也就是說，必須將它們的結構序列化為一個VariantTensorData對象，類結構如下：

struct VariantTensorData {
    string type_name;
    string metadata;
    std::vector<Tensor> tensors;
};

如果對象內包含指向其它張量的引用，可以把它們包含在tensors中，把其它的元數據內容放入metadata中。

versions

提供了版本信息檢查的功能。

tensorflow源碼解析之framework拾遺

節點 repeat 普通相互 eve pen lib ash 除了把framework中剩余的內容，按照文件名進行了簡單解析。時間原因寫的很倉促，算是占個坑，後面有了新的理解再來補充。 allocation_description.proto 一個對單次內存分配結果進行

tensorflow源碼解析之framework拾遺

allocation_description.proto

attr_value

bfloat16

cancellation

common_shape_fns

control_flow

cost_graph

fake_input

load_library

log_memory

memory_types

numeric_op

numeric_types

register_types

register_types_traits

rendezvous

session_state

step_stats

summary.proto

type_index

types

type_traits

unique_tensor_references

variant

versions

相關推薦