Caffe2原始碼解析

摘要： 寫在前面 上一篇文章對Caffe2中的core模組進行了簡單拆解Caffe2原始碼解析之core，本篇給出其它模組的拆解，目的是大致瞭解每個模組的內容和目標，進一步理解Caffe2的整體框架。內容不多，略做整理如下。 目錄 core proto...

寫在前面

上一篇文章對Caffe2中的core模組進行了簡單拆解Caffe2原始碼解析之core，本篇給出其它模組的拆解，目的是大致瞭解每個模組的內容和目標，進一步理解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

core

參見ofollow,noindex" target="_blank">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

先來看下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 ，歡迎大家點星星。