Caffe2源碼解析

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

遍歷多線程準備 var on() 大致編譯期 pack eset

寫在前面

上一篇文章對Caffe2中的core模塊進行了簡單拆解Caffe2源碼解析之core，本篇給出其它模塊的拆解，目的是大致了解每個模塊的內容和目標，進一步理解Caffe2的整體框架。內容不多，略做整理如下。

core

參見Caffe2源碼解析之core

proto

包含了Caffe2中常用的protobuf定義，非常重要。我們按照所在文件進行介紹

caffe2.proto

首先是TensorProto，它表示張量序列化後的結果，包括了張量的維度、數據類型、數據值、名稱、所在設備等信息，如下：

message TensorProto {
    repeated int64 dims = 1;
    optional DataType data_type = 2 [default = FLOAT];
    repeated float float_data = 3 [packed = true];
    repeated int32 int32_data = 4 [packed = true];
    optional bytes byte_data = 5;
    repeated bytes string_data = 6;
    repeated double double_data = 9 [packed = true];
    repeated int64 int64_data = 10 [packed = true];
    optional string name = 7;
    
    //張量序列化之前所在的設備
    optional DeviceOption device_detail = 8;
    //張量在chunks中的位置
    message Segment {
        required int64 begin = 1;
        required int64 end = 2;
    }
    optional Segment segment = 11;
}

在core模塊中講到，Caffe2為了支持低精度的模型訓練，設計了qtensor，當時沒有詳細介紹它的本質，實際上qtensor是對原張量進行了歸一化，即減去bias再除以scale，然後對結果進行低精度表示，節省存儲空間。因此在qtensor的序列化結果中，需要對歸一化的參數進行記錄，如下：

message QTensorProto {
    repeated int64 dims = 1;
    required int32 precision = 2;
    required double scale = 3;
    required double bias = 4;
    required bool is_signed = 5;
    repeated int32 data = 6 [packed = true];
    optional string name = 7;
    optional TensorProto.DataType data_type = 8 [default = INT32];
}

對於多個Tensor，可以使用TensorProto的復數版本，TensorProtos來存儲。當然這只針對較小的張量，如果張量比較大，建議使用DB存儲。

對於張量的形狀，也有一個結構來表示，TensorShape。記得在Tensorflow中，對於張量形狀的某些維度，在運行前可能並不是完全知道，因此這裏在TensorShape的定義中，會添加參數對未知張量維度做處理。

message TensorShape {
    repeated int64 dims = 1;
    optional TensorProto.DataType data_type = 2 [default = FLOAT];
    repeated int32 unknown_dims = 3;
    optional bool unknown_shape = 4 [default = false];
    optional string name = 5;
}

參數用於對操作的描述（詳見下文的OperatorDef定義），一個命名的參數要麽包含單個的浮點型、整型或者字符串數據，要麽包含上述類型的數組，如下：

message Argument {
    optional string name = 1;
    optional float f = 2;
    optional int64 i = 3;
    optional bytes s = 4;
    optional NetDef n = 8;
    repeated float floats = 5;
    repeated int64 ints = 6;
    repeated bytes strings = 7;
    repeated NetDef nets = 9;
}

目前Caffe2支持的設備類型：

enum DeviceType {
    CPU = 0;
    CUDA = 1;
    MKLDNN = 2;
    OPENGL = 3;
    OPENCL = 4;
    IDEEP = 5;
    HIP = 6;
    COMPILE_TIME_MAX_DEVICE_TYPES = 7;
    ONLY_FOR_TEST = 20901701;
}

目前Caffe2對於不同設備的描述proto還都是一致的，如果某個設備沒有包含其中的某個字段，那麽這個字段將被忽略。

message DeviceOption {
    optional int32 device_type = 1 [default = 0]; //0 is CPU
    optional int32 cuda_gpu_id = 2;
    optional uint32 random_seed = 3;
    optional string node_name = 4;
    optional int32 numa_node_id = 5;
    repeated string extra_info = 6;
    optional int32 hip_gpu_id = 7;
}

接下來是操作的定義：

message OperatorDef {
    repeated string input = 1; //輸入的blob名稱
    repeated string output = 2; //輸出的blob名稱
    optional string name = 3;
    optional string type = 4; //操作的類型，從操作註冊器中創建操作對象時，需要這個信息
    optional string type = 4;
    repeated Argument arg = 5;
    optional DeviceOption device_option = 6; //操作運行所需要的設備
    
    //對於當前操作來說，如果對於指定的運行設備有多個計算引擎，這裏可以指定一個具體的實現引擎。如果用戶指定了一個引擎，但這個引擎在Caffe2的二進制包中並不存在，那麽使用默認引擎
    optional string engine = 7;
    
    //控制輸入，與Tensorflow中的控制輸入類似，表達運行的先後順序，而不是數據的輸入。它僅在調度時被Net類使用
    repeated string control_input = 8;
    
    //is_gradient_op參數僅在形狀推斷（shape inference，與Tensorflow中類似）時使用，沒有運行時的作用
    optional bool is_gradient_op = 9 [default = false];
    optional string debug_info = 10;
}

接下來NetDef的定義：

message NetDef {
    optional string name = 1;
    repeated OperatorDef op = 2;
    optional string type = 3; //network的執行方式，默認是simple
    optional DeviceOption device_option = 5; //整個net上所有操作的設備信息，在這裏設置可以避免給每個操作單獨設置
    repeated Argument arg = 6; //參數，包括num_workers，即當圖被並行執行的時候，worker的數量
    
    repeated string external_input = 7;
    repeated string external_output = 8;
}

Caffe2中也可以像Tensorflow那樣進行叠代計算，它使用了一個結構叫做ExecutionStep，如下：

message ExecutionStep {
    optional string name = 1;
    
    //ExecutionStep要麽可以包含一個substep的集合，要麽可以包含一些要運行的network的名稱，但兩者不能同時被設置
    repeated ExecutionStep substep = 2;
    repeated string network = 3;
    
    //當前的叠代需要運行的輪次，substeps和networks需要被順序執行，每次執行被視為一輪叠代
    optional int64 num_iter = 4;
    
    //叠代執行結束的判斷條件
    optional string criteria_network = 5;
    
    //如果這個字段被設置，那麽就周期性的執行
    optional int64 run_every_ms = 11;
    
    //對於sub-steps，是順序執行還是並行執行
    optional bool concurrent_substeps = 6;
    
    //一個用來判斷當前執行是否需要終結的標誌
    optional string should_stop_blob = 9;
    
    //如果為真，則當前執行僅執行一次，註意僅當should_stop_blob有效時才有效
    optional bool only_once = 10;
    
    //是否為當前執行構建一個子workspace
    optional bool create_workspace = 12;
    
    //子執行的並行度
    optional int32 num_concurrent_instances = 13;
}

如果說一個ExecutionStep是一次叠代執行，那麽Plan就是一個完整的執行計劃，後者包含前者：

message PlanDef {
    optional string name = 1;
    repeated NetDef netowrk = 2;
    repeated ExecutionStep execution_step = 3;
}

對於那些內部並不是Tensor的Blob，Caffe2定義了如下的結構：

message BlobProto {
    optional string name = 1;
    optional string type = 2;
    optional TensorProto tensor = 3;
    optional bytes content = 4;
    optional QTensorProto qtensor = 5;
    optional int32 content_num_chunks = 6;
    optional int32 content_chunk_id = 7;
}

最後，是對DBReader進行序列化的對象：

message DBReaderProto {
    optional string name = 1;
    optional string source = 2;
    optional string db_type = 3;
    optional string key = 4;
}

hsm.proto

Word2Vec是早年Google提出的一個模型，目的是根據語料庫獲得詞嵌入（embedding）。其中為了提高訓練的速度提出了兩種技術，一種是負采樣（Negative Sampling），另外一種就是Hierarchical Softmax。因此，Caffe2專門設計了一個HSM操作，這個文件裏包含的就是與之相關的proto，我們僅給出proto名稱，含義比較顯然：

message NodeProto;
message TreeProto;
message HierarchyProto;
message PathProto;
message PathNodeProto;

metanet.proto

MetaNetDef，顧名思義，包含了NetDef的元數據。其結構如下：

message MetaNetDef {
    repeated BlobMap blobs = 1;
    repeated NetsMap nets = 2;
    optional ModelInfo modelInfo = 3;
    repeated PlanMap plans = 4;
    repeated StringMap applicationSpecificInfo = 5;
}

其中，對應的xxMap結構很簡單，都是鍵值對，ModelInfo相對比較復雜，我們看下詳細的定義：

message ModelInfo {
    optional string project = 1;
    optional string modelClass = 2;
    optional string version = 3;
    optional string predictorTtype = 4;
    optional string modelId = 5;
}

cuda_rtc

cuda核生成相關的輔助代碼。

db

在Caffe2的執行過程中，需要重復使用和共享的參數，會被記錄在一個db當中。在core模塊中我們介紹過，db就是一個kv存儲，這裏包含了4種Caffe2中會用到的db，如下：

graph TB db-->|派生|LevelDB db-->|派生|LMDB db-->|派生|ProtoDB db-->|派生|ZmqDB

distributed

Caffe2的分布式實現，依賴外部存儲來保存共享的參數。常用的外部存儲包括文件和redis。

外部存儲的句柄用StoreHandler來表示，它包含了以下的核心API：

class StoreHandler {
  public:
    virtual void set(...) = 0;
    virtual std::string get(...) = 0;
    virtual int64_t add(...) = 0;
    virtual bool check(...) = 0;
    virtual void wait(...) = 0;
};

對應到計算圖中，就有4個對store操作的op與之對應，如下：

graph TB Operator-->|派生|StoreSetOp Operator-->|派生|StoreGetOp Operator-->|派生|StoreAddOp Operator-->|派生|StoreWaitOp

剛才提到了，常用的存儲方式為文件存儲和redis存儲，對應有兩種存儲句柄：

graph TB StoreHandler-->|派生|RedisStoreHandler StoreHandler-->|派生|FileStoreHandler

另外，還有兩個創建存儲的操作，如下：

graph TB Operator-->|派生|FileStoreHandlerCreateOp Operator-->|派生|RedisStoreHandler

ideep

目前還不清楚具體含義。

image

關於圖像的操作，其中最重要的是對於圖像讀取的操作，ImageInputOp，它繼承自PrefetchOperator，包含了圖像讀取的一系列功能。

mkl

MKL全稱是Intel Math Kernel Library，是英特爾提供的數學核心庫，它對大量的數學過程進行了處理器級別的優化。這裏包括了MKL相關的操作定義。註意，Tensorflow中也用到了MKL去優化數學運算，只不過它是在圖優化的過程中，將MKL作為一種圖優化遍歷被引入，而Caffe2中將MKL直接融入到了操作內部。

mobile

針對移動平臺的特殊處理，具體還沒看。

mpi

Caffe2中的分布式計算，通過mpi實現。mpi的核心作用是在不同機器上的分布式進程中，進行數據傳輸和消息同步。針對mpi中的核心操作，比如Broadcast，Reduce等，Caffe2都給出了對應的操作來執行，具體如下：

graph TB Operator-->|派生|MPICreateCommonWorldOp Operator-->|派生|MPIBroadcastOp Operator-->|派生|MPIReduceOp Operator-->|派生|MPIAllgatherOp Operator-->|派生|MPIAllreduceOp Operator-->|派生|MPISendTensorOp Operator-->|派生|MPIReceiveTensorOp

observers

給出了4種不同觀察器的定義，如下：

operator_attaching_net_observer，負責給net中的每一個operator添加觀察器；
profile_observer，負責對每個操作或整張圖的執行消耗進行觀察；
runcnt_observer，負責對每個操作或者整張圖的運行次數進行觀察；
time_observer，負責對每個操作或者整張圖的運行時間進行觀察；

onnx

目前還不清楚。

operators

操作的具體定義放在這裏，代碼量巨大，沒來得及細看。

opt

優化相關的類和函數，與Tensorflow一樣，Caffe2也是通過對圖遍歷的方式實施優化，所有的優化遍歷類必須繼承自OptimizationPass，它的具體定義如下：

class OptimizationPass {
  public:
    OptimizationPass(NNModule* nn) : nn_(nn) {}
    virtual void run() = 0;
    virtual ~OptimizationPass(){}
    
  protected:
    NNModule* nn_;
};

perfkernels

性能優化相關的kernel。

predictor

一個predictor就是一個參數都確定好了的net。在深度學習中，我們通常會把待學習的模型表示為net，然後通過叠代的圖計算，確定模型參數，將net轉換為predictor。下面我們看下predictor的結構：

class Predictor {
  public:
    Predictor(const NetDef& init_net, const NetDef& run_net, Workspace* parent = nullptr, bool run_init = true, int optimization = 1);
    Predictor(PredictorConfig config);
    
    //以下是對()的重載，給定輸入得到輸出
    bool operator()(const TensorMap& inputs, TensorList* outputs);
    bool operator()(const TensorMap& inputs, TensorList* outputs);
    bool operator()(const TensorMap& inputs, TensorMap* outputs);
    
    const NetDef& def() const {
        return *config_.predict_net;
    };
    
    Workspace* ws(){
        return config_.ws.get();
    };
    const std::vector<std::string>& input_names() const {
        return config_.input_names;
    }
    const std::vector<std::string>& output_names() const {
        return config_.output_names;
    }
  private:
    bool run_map_workspace(const TensorMap& inputs);
    PredictorConfig config_;
};

其中，Predictor類最重要的一個私有數據成員是config_，我們看下PredictorConfig的定義：

struct PredictorConfig {
    std::shared_ptr<PredictorParameters> parameters;
    std::shared_ptr<NetDef> predict_net;
    std::vector<std::string> input_names;
    std::vector<std::string> output_names;
    std::vector<std::string> parameter_names;
    std::shared_ptr<Workspace> ws;
};

queue

與Tensorflow類似，Caffe2也利用隊列對多個線程進行同步，比如在多線程讀取輸入數據的時候。對隊列的所有動作都必須通過“操作”來完成，因此Caffe2又定義了隊列相關的操作。

先來看下BlobsQueue的定義：

class BlobsQueue : public std::enable_shared_from_this<BlobsQueue> {
  public:
    bool blockingRead(...);
    bool blockingWrite(...);
    void close();
  private:
    size_t numBlobs_;
    std::mutex mutex_;
    std::condition_variable cv_;
    int64_t reader_{0};
    int64_t writer_{0};
    std::vector<std::vector<Blob*>> queue_; //核心隊列數據
    const std::string name_;
};

註意看其中的數據元素queue_，它就是BlobsQueue的核心隊列數據。

另外，BlobsQueue，也可以被看做是一種db，因此Caffe2定義了BlobQueueDB：

class BlobsQueueDB : public DB {
  public:
    BlobsQueueDB(...);
    void Close() override {}
    unique_ptr<Cursor> NetCursor() override{...}
    unique_ptr<Transaction> NewTransaction() override {...}
  private:
    std::shared_ptr<BlobsQueue> queue_;
    int key_blob_index_;
    int value_blob_index_;
    float timeout_secs_;
};

另外，Caffe2還針對BlobsQueue提出了提出了對隊列進行處理的“操作”，把常用的隊列處理方式，如入隊、出隊等，抽象為操作：

graph TB Operator-->|派生|CreateBlobsQueueOp Operator-->|派生|EnqueueBlobsOp Operator-->|派生|DequeueBlobsOp Operator-->|派生|CloseBlobsQueueOp Operator-->|派生|SafeEnqueueBlobsOp Operator-->|派生|SafeDequeueBlobsOp Operator-->|派生|WeightedSampleDequeueBlobsOp

另外，為了能支持一次多數據入隊，Caffe2設計了RebatchingQueue類，它的簡要結構如下：

class RebatchingQueue {
  public:
    bool enqueueOne(...);
    bool enqueueMany(...);
    bool dequeue(...);
  private:
    std::vector<std::vector<TensorCPU>> queue_;
};

與BlobsQueue最大的區別有兩點，第一，核心數據queue_中存儲的是TensorCPU而不是Blob*，第二，擁有EnqueueOne和EnqueueMany兩種入隊操作。

與BlobsQueue類似，Caffe2也為RebatchingQueue準備了對其進行處理的“操作”，與BlobsQueue類似，這裏不再贅述。

sgd

包含了與隨機梯度下降有關的操作。基本上可以根據文件名猜測含義，這裏僅列出文件名前綴，感興趣的讀者可以查閱源碼：

adadelta_op
adagrad_op
adam_op
clip_tensor_op
fp16_momentum_sgd_op
fp32_momentum_sgd_op
ftrl_op
gftrl_op
iter_op
lars_op
learning_rate_adaption_op
learning_rate_functors
learning_rate_op
momentum_sgd_op
rmsprop_op
wngrad_op
yellowfin_op

有機會可以仔細研讀下其中的細節。

transform

根據core模塊的內容我們知道，這裏包含的是對圖進行變換的方法。主要包括4種：

//公共子項消除，CSE，與Tensorflow類似
common_subexpression_elimination

//對卷積操作進行變換，提高效率
conv_to_nnpack_transform

//模式替換，允許你使用簡單的接口定義模式替換規則，只需定義一模式子圖和一個替換子圖，在原圖中尋找模式子圖，然後替換為替換子圖即可
pattern_net_transform

//單個操作的原地轉換
single_op_transform

這些類形成了如下的繼承體系：

graph TB Transform-->|派生|CommonSubexpressionEliminationTransform Transform-->|派生|SingleOpTransform Transform-->|派生|PatternNetTransform SingleOpTransform-->|派生|ConvToNNPackTransform

util

應用類和函數，比較瑣碎，暫時沒有細看。

python

通過前面的介紹我們了解到，Caffe2的核心代碼是用"C++"實現的，為了方便在python中進行調用，需要一個工具，幫助python調用"C++"代碼。這樣的工具有很多，比如boost.python, swig，ctypes，pybind11等。Caffe2選擇了pybind11，因為它對"C++"11支持的比較好，而且API比較簡單。而Tensorflow中python前端調用"C++"後端使用的是swig，其實swig對"C++"11也能支持。兩種設計選擇的優劣目前的知識我們還不好評判。

具體的接口文件，是_import_c_extention.py，它首先會嘗試載入gpu版本的Caffe2後端，如果失敗了，會嘗試載入CPU版本。其中，對於CPU後端的導入是通過如下的語句：

from caffe2.python.caffe2_pybind11_state import *

因此，在編譯完成後，caffe2/python目錄下會生成一個名為caffe2_pybind11_state.so的文件，是包含了Caffe2的"C++"後端的動態鏈接庫，可以被python載入。

contrib

同Tensorflow的contrib文件夾一樣，包含了第三方貢獻的、未正式加入Caffe2的模塊，這裏面大部分代碼是用python開發的。隨著版本叠代，經測試穩定後，這些模塊會逐漸加入Caffe2的python模塊。

寫在後面

看過Tensorflow和Caffe2的核心代碼之後，講一講自己的感受。

代碼模塊性，Tensorflow代碼的模塊性做的非常好，基礎框架、運行時、圖表示、圖優化、op、kernel都區分的清清楚楚，而Caffe2的代碼顯得有些混雜，操作到處都是，給代碼閱讀帶來了一點障礙。
代碼規範性，Tensorflow代碼的規範性要好很多，雖然核心代碼是多個作者完成的，但代碼風格非常統一，文件頭的協議也非常一致。反觀Caffe2的代碼，協議混亂，代碼風格不統一，東拼西湊的感覺比較強烈，代碼在形式上的美感不足。
架構合理性，Tensorflow的野心很大，它的終極目標是變成一個全新的、面向數據流圖計算的編程語言。這種編程語言基於op原語，利用op和kernel將編譯期和運行期明確的區分開來，同時，它對於同一個數據流圖的多線程並行執行機制，也像極了CPU流水線處理機制，因此，應該說，深度神經網絡只是Tensorflow的一個副產品，它的真實價值遠不止於此。反觀Caffe2，很多設計有些短視了（比如用redis為中介做分布式執行），在提供更多靈活性的同時，也限制了它的高度。

當然，以上只是個人的一些猜測，隨著理解的深入，我也會及時回來修正自己的觀點，也歡迎大家來討論。

最後，我在github上新建了一個repo，pytorch_notes，歡迎大家點星星。

Caffe2源碼解析

Caffe2源碼解析

寫在前面

目錄

core

proto

caffe2.proto

hsm.proto

metanet.proto

cuda_rtc

db

distributed

ideep

image

mkl

mobile

mpi

observers

onnx

operators

opt

perfkernels

predictor

queue

sgd

transform

util

python

contrib

寫在後面

相關推薦