寫在前面

在對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

記憶體分配器，包含了CPU和GPU兩種。

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註冊的函式；