1. 程式人生 > >Caffe2源碼解析之core

Caffe2源碼解析之core

numeric pua 基本 參數 set 記錄 4.2 oca www

寫在前面

在對Tensorflow的後端源碼進行了拆解(參見tensorflow源碼解析系列文章索引)之後,很想跟其它深度學習框架的實現進行對比,根據框架的流行程度,先選擇了Pytorch。Pytorch的後端核心是直接復用了Caffe2,因此本文針對Caffe2源碼的core模塊進行了簡單拆解。

目錄

  • 數據存儲與表示
    • storage
    • tensor
    • blob
    • qtensor
  • 操作
    • observer observable
    • operator
    • 操作求導
    • operator_schema
    • context
  • 計算圖
    • graph
    • net
    • transform
  • 運行時
    • allocator
    • db
    • registry
    • module
    • scope_guard
    • workspace
    • init

1. 數據存儲與表示

1.1 storage

Caffe2中對數據存儲的最底層的描述是Storage,它實際上是指向StorageImpl的共享指針,後者包含數據類型、數據指針、容量、數據所在設備等信息。Storage的定義如下:

using Storage = std::shared_ptr<StorageImpl>;
class StorageImpl {
  public:
    //...
  protected:
    using DataPtr = std::shared_ptr<void>;
    int64_t capacity_ = 0;
    DataType data_type_;
    DataPtr data_ptr_;
    DeviceType device_type_ = CPU;
};

1.2 tensor

Caffe2中的數據統一使用Tensor表示,Tensor由TensorImpl實現,後者包含一個Storage。

graph LR Tensor-->|包含|TensorImpl TensorImpl-->|包含|Storage Storage-->|指向|StorageImpl

TensorImpl的定義如下:

class TensorImpl {
  public:
    //...
  protected:
    using DimVector = std::vector<TIndex>;
    DimVector dims_; //張量的維度
    TIndex size_ = -1; //張量中包含的元素數量
    Storage storage_; //底層存儲
};

Tensor並非繼承自TensorImpl,而是在內部包含了一個指向TensorImpl的指針,如下:

class Tensor final {
  protected:
    using TensorImplPtr = c10::intrusive_ptr<TensorImpl, UndefinedTensorImpl>;
    TensorImplPtr impl_;
  //...
};

對Tensor的方法調用,通過重定向給TensorImpl實現。

1.3 blob

Blob是一個容器,包含了一個指針和這個指針指向內存的數據類型,在Caffe2中,大部分情況下Blob都包含一個指向Tensor的指針。

class Blob {
  public:
    //...
  private:
    TypeMeta meta_;
    void* pointer_ = nullptr;
    DestroyCall destroy_ = nullptr;
};

為了方便對Blob進行傳輸,定義了其序列化和反序列化的類,分別是BlobSerializerBase和BlobDeserializerBase,以及對應的為Tensor準備的序列化和反序列化類。

graph TB BlobSerializerBase-->|派生|TensorSerializer BlobDeserializerBase-->|派生|TensorDeserializer

1.4 qtensor

低精度的張量,為了便於快速進行低精度的整數乘法計算。具體的做法是,用更低的位數來表示整數,比如,用3個bit表示無符號整數,用4個bit表示有符號整數。低精度張量可以在略微損失模型精度的情況下,大大降低計算復雜度和存儲空間大小。

操作

2.1 Observer Observable

Caffe2使用ObserverBase和Observable兩個類實現了觀察者模式。ObserverBase是基礎觀察器,用戶可以通過繼承此類創建新的觀察器,而Observable是可被觀察屬性,用戶可以通過繼承此類獲得可觀察屬性。

ObserverBase提供了觀察器的統一接口,比較簡單:

class ObserverBase {
  public:
    virtual void Start() {}
    virtual void Stop() {}
    T* subject() const {
        return subject_;
    }
  protected:
    T* subject_;
};

其中,subject_表示被觀察對象的指針。

Observable封裝了可被觀察屬性,內部包含了一個觀察器的列表,結構如下:

class Observable {
  public:
    using Observer = ObserverBase<T>;
    const Observer* AttachObserver(std::unique_ptr<Observer> observer){} //添加觀察器
    std::unique_ptr<Observer> DetachObserver(const Observer* observer_ptr){} //解除觀察器
    virtual size_t NumObservers() {
        return num_observers_;
    } //觀察器的數量
    void StartAllObservers(){} //啟動所有觀察器
    void StopAllObservers(){} //關閉所有觀察器
  private:
    Observer* observer_cache_;
    size_t num_observers_ = 0;
  protected:
    std::vector<std::unique_ptr<Observer>> observer_list_; //觀察器列表
};

2.2 Operator

Operator代表操作的具體實現,相當於Tensorflow中的kernel。Operator繼承自OperatorBase,而後者繼承自Observable,所以在Caffe2中,“操作”本質上是一個可觀察的對象。

graph LR Observable-->|派生|OperatorBase OperatorBase-->|派生|Operator

OperatorBase類包含了操作需要的基本數據元素和接口:

class OperatorBase {
  private:
    Workspace* operator_ws_;
    std::shared_ptr<const OperatorDef> operator_def_;
    DeviceOption device_option_;
    std::string engine_;
    std::string type_;
    vector<const Blob*> inputs_;
    vector<Blob*> outputs_;
};

OperatorBase中包含了輸入和輸出的內存指針,可見,在Caffe2中,Operator本質上是一個運行時的對象,這與Tensorflow中Op的設計理念不同,在Tensorflow中,Op是一個編譯時對象,僅規定了操作的類型和目標,並不包含具體數據,具體的計算實際上是通過Kernel完成的。

Operator繼承自OperatorBase類:

class Operator : public OperatorBase {
  public:
    bool Run(int stream_id = 0) final {...}
    bool RunAsync(int stream_id = 0) final {...}
    virtual bool RunOnDevice() = 0;
};

實際上,Run和RunAsync最終都調用了RunOnDevice,完成實際的計算。

如果我們需要使用一些c10中定義的操作,需要將其轉換為在Caffe2中可以調用的操作,可以通過如下的宏進行轉換:

REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH(C10Add, C2MyAddOpName)

上述例子中,我們把一個C10Add操作,包裝成C2MyAddOpName操作,供我們使用。為了實現這個功能,Caffe2還提供了一個包裝類,C10OperatorWrapper。

2.3 操作求導

為了對操作求導,Caffe2推出了一個導數操作生成類,GradientMakerBase,方便用戶定義對於某個操作的導數。類包含的數據成員如下:

//為密集和稀疏的blob提供統一的接口
struct GradientWrapper {
    string dense_;
    string indices_;
    string values_;
    inline bool IsDense(){}
    inline bool IsSparse(){}
    inline bool IsEmpty(){}
};
class GradientMakerBase {
  protected:
    const OperatorDef& def_;
    const vector<GradientWrapper>& g_output_;
    vector<GradientWrapper> g_input_;
};

可見,GradientMakerBase僅提供了輸入輸出,以及原操作。用戶可以根據原操作,定制導數。

2.3 operator_schema

OpSchema是對操作的靜態描述,相當於Tensorflow中的Op,包含的信息如下:

class OpSchema {
  private:
    string type_;
    string file_;
    string doc_;
    string onnx_schema_;
    std::vector<Argument> args_{};
    std::vector<std::pair<const char*, const char*>> input_desc_{};
    std::vector<std::pair<const char*, const char*>> output_desc_{};
    int line_ = 0;
    int min_input_ = 0;
    int max_input_ = std::numeric_limits<int>::max();
    int min_output_ = 0;
    int max_output_ = std::numeric_limits<int>::max();
    bool private_ = false;
    bool inputs_can_cross_devices_ = false;
    std::function<bool(int)> num_inputs_allowed = [](int) { return true; }
    std::function<bool(int)> num_outputs_allowed = [](int) { return true; }
    std::function<bool(int,int)> num_inputs_outputs_allowed_ = [](int,int) { return true; }
    std::function<int(int)> calculate_output_;
    std::function<bool(int,int)> inplace_allowed_ = [](int,int){}
    std::function<bool(int,int)> inplace_enforced_ = [](int,int){}
    TensorInferenceFunctionType tensor_inference_function_ = {...}
    std::unique_ptr<CostInferenceFunctionType> cost_inference_function_ = nullptr;
    DeviceInferenceFunctionType device_inference_function_ = {...}
};

另外Caffe2也提供了一個對於OpSchema的註冊類OpSchemaRegistry,如下:

class OpSchemaRegistry {
  private:
    static CaffeMap<string, OpSchema>& map();
};

2.4 context

Caffe2中的context,其實就是Tensorflow中的OpKernelContext,為操作的實際計算提供通用的支持,主要包含內存拷貝的接口。所有實際的Context類必須繼承自BaseContext,而Caffe2為我們準備了一個標準的Context接口,CPUContext類。另外,也同樣為GPU準備了一個CUDAContext類。

graph LR BaseContext-->|派生|CPUContext BaseContext-->|派生|CUDAContext

3. 計算圖

3.1 graph

Graph表示圖的結構,圖包含節點,節點包含操作。

graph LR Graph-->|包含|Node Node-->|包含|OperatorDef

Node包含的數據成員:

class Node {
  public:
    OperatorDef op;
    bool active = true; //操作是否被transformation刪除
    std::map<int, std::vector<string>> parents;
    std::vector<int, std::vector<string>> children;
}

Graph包含的私有數據成員:

class Graph {
  private:
    NetDef netdef_;
    std::set<string> external_input_;
    std::set<string> external_output_;
    std::vector<Node> nodes_;
}

3.2 net

Net是一個可運行的Graph,包含了一個圖的所有“操作”,以及它們的上下文。它繼承自Observable,本質上是一個可觀察的對象。數據成員如下:

class NetBase : public Observable<NetBase>{
  public:
    virtual bool Run(){...}
    virtual bool RunAsync();
  protected:
    vector<string> external_input_;
    vector<string> external_output_;
    string name_;
    vector<const Event*> events_;
    std::shared_ptr<const NetDef> net_def_;
};

NetBase派生出了三種子類,第一種是AsyncNetBase,它包含了異步執行網絡所必須的數據和接口:

class AsyncNetBase : public NetBase {
  public:
    bool RunAsync() override;
  protected:
    bool canSchedule(...);
    std::vector<OperatorBase*> operators_;
    std::vector<dag_utils::OperatorNode> operator_nodes_;
    std::vector<std::vector<int>> chains_;
    std::vector<dag_utils::OpGraphNode> chain_nodes_;
    dag_utils::ExecutionChains execution_chains_;
};

第二種是SimpleNet,它表示了一種對圖的單線程的順序執行模式。 第三種是DAGNetBase,它表示了一種對圖的多線程的dag執行模式。 相關的net類形成了一個繼承體系:

graph TB Observable-->|派生|NetBase NetBase-->|派生|AsyncNetBase AsyncNetBase-->|派生|AsyncSchedulingNet NetBase-->|派生|DAGNetBase DAGNetBase-->|派生|DAGNet NetBase-->|派生|SimpleNet DAGNetBase-->|派生|AsyncDAGNet AsyncNetBase-->|派生|AsyncPollingNet

3.3 transform

transform是一種針對Caffe2的NetDef結構的操作,它將NetDef作為輸入,輸出新的經過變換的NetDef。它的工作步驟包括:

  • 從舊的NetDef中構建一張圖,這張圖中保存了節點的連接信息;
  • 在圖中匹配指定的模式,找到它想要更改的子圖;
  • 用新的操作替換匹配到的子圖;
  • 根據圖構建一個新的NetDef並返回;

Transform功能的實現,依賴於三個功能函數,如下:

  • PatternRule(模式規則),它決定了對於一張子圖和一個節點,是否可以將這個節點加入這個子圖中;
  • ValidatorRule(驗證規則),它決定了一張子圖是否是匹配的;
  • ReplaceRule(替換規則),它對一張匹配的子圖進行替換;

常用的模式如下:

  • CONNECTED_SUBGRAPH,連接子圖,它只能匹配連接的子圖。比如對於圖(1)-->(2)-->(3)-->(4),它能夠匹配到[2,3]和[4,3],但不能匹配到[2,4];
  • SORTED_WRT_EXECUTION_ORDER,執行序模式,它只能匹配符合執行順序的子圖,節點之間不一定需要有連接,它比General模式要快,例如對於圖(1)-->(2)-->(3)-->(4),它可以匹配到[2,4],[3,4],但不能匹配到[3,1],[4,3];
  • GENERAL,它可以匹配到任何子圖,比如,對於圖(1)-->(2)-->(3)-->(4)來說,它可以匹配到子圖[2,4],[3,4],[4,2,1]等;

4. 運行時

4.1 allocator

內存分配器。

graph TB CPUAllocator-->|派生|DefaultCPUAllocator CPUAllocator-->|派生|PinnedCPUAllocator

4.2 db

DB類是對kv存儲的抽象。包含了用於讀取DB數據的Cursor類,用於寫DB數據的Transaction類,DB讀取的包裹類DBReader,對DBReader進行序列化和反序列化的DBReaderSerializer和DBReaderDeserializer類。

graph TB DB-->|讀數據時的遊標類|Cursor DB-->|寫數據時的事務類|Transaction DB-->|讀數據包裝|DBReader DBReader-->|序列化|DBReaderSerilizer DBReader-->|反序列化|DBReaderDeserilizer

4.3 registry

註冊類,key為字符串,value可以為任意的類。結構如下:

class Registry {
  private:
    CaffeMap<SrcType, Creator> registry_;
    CaffeMap<SrcType, string> help_message_;
};

4.4 module

查看Caffe2已載入的模塊,以及載入指定模塊。模塊指的是動態鏈接庫。

4.5 scope_guard

是“初始化即資源獲取”原語的實現,它保證了,如果不顯式說明,函數的執行就會離開當前的scope。

4.6 workspace

Workspace包含了所有的運行時對象,包括blob和net,它是所有這些對象的擁有者,負責對這些對象進行管理。

class Workspace {
  private:
    typedef CaffeMap<string, unique_ptr<Blob>> BlobMap;
    BlobMap blob_map_;
    typedef CaffeMap<string, unique_ptr<NetBase>> NetMap;
    NetMap net_map_;
    const string root_folder_;
    const Workspace* shared_;
    std::unordered_map<string, std::pair<const Workspace*, string>> forwarded_blobs_;
    std::unique_ptr<ThreadPool> thread_pool_;
    std::mutex thread_pool_creation_mutex_;
    std::shared_ptr<Bookkeeper> bookkeeper_;
};

4.7 init

初始化整個Caffe2的運行環境,運行機制是,把需要在環境初始化中運行的函數註冊到註冊器中,初始化時,會在不同時期運行不同註冊器中的函數。核心的函數如下:

CAFFE2_API bool GlobalInit(int* pargc, char*** argv);

整個初始化過程分為三步:

  • 先運行通過REGISTER_CAFFE2_EARLY_INIT_FUNCTION註冊的函數;
  • 再解析Caffe的命令行參數,並啟動日誌記錄系統;
  • 最後運行通過REGISTER_CAFFE2_INIT_FUNCTION註冊的函數;

Caffe2源碼解析之core