mxnet原始碼閱讀筆記之include
寫在前面
mxnet程式碼的規範性比Caffe2要好,看起來核心程式碼量也小很多,但由於對dmlc其它庫的依賴太強,程式碼的獨立性並不好。依賴的第三方庫包括:
cub dlpack dmlc-core googletest mkldnn mshadow onnx-tensorrt openmp ps-lite tvm
如果對於這些第三方庫沒有足夠的理解,mxnet的核心程式碼看起來比較費勁。因此時間原因,本篇僅解析了mxnet對外的介面include目錄,並且對於嚴重依賴第三方庫的檔案沒有深入探究,只能算作一篇不完整的原始碼閱讀筆記了。後續有時間的話,再回來迭代。
目錄
- storage
- tensor_blob
- ndarray
- resource
- kvstore
- base
- operator
- engine
- executor
- rtc
- graph_attr_types
- op_attr_types
- imperative
- operator_util
- c_api
storage
Storage是一個跨裝置的記憶體管理類,它提供了記憶體分配和回收的功能,但並不儲存分配的記憶體,真正的記憶體指標分配在Storage類內部的Handle結構體中:
struct Handle { void * dptr{nullptr}; //記憶體地址 size_t size{0}; Context ctx; int shared_pid{-1}; int shared_id{-1}; }; class Storage { public: Handle Alloc(size_t size, Context ctx) {...}; virtual void Alloc(Handle* handle) = 0; virtual void Free(Handle handle) = 0; };
tensor_blob
TBlob類可以表示任意維度、在任意裝置上、任意資料型別的張量,它是NDArray的內部儲存,是mxnet中最底層的資料結構。但本質上它是對DLTensor的代理,DLTensor定義在第三方庫dlpack中的dlpack.h檔案中,以下是它們的關係:
graph LR NDArray-->|包含|TBlob TBlob-->|包含|DLTensor
ndarray
ndarray是mxnet中的核心資料結構,代表了多維資料,類似於Tensorflow中的Tensor。本質上它借鑑了numpy中關於ndarray的定義,一部分ndarray是包含實際資料的,另外一些ndarray並不包含實際資料,它們只是其他ndarray的檢視。舉例說明,ndarrayA是一個[1x12]的多維陣列,儲存了12個元素,ndarrayB是一個[3x4]的多維陣列,它底層的資料由ndarrayA提供,因此A和B共享了記憶體,B僅是A的一個檢視。
ndarray內部由chunk結構提供實際的資料儲存,先來看下chunk:
struct Chunk { Storage::Handle shandle; std::vector<Storage::Handle> aux_handles; bool static_data; //如果為真,表示該資料是靜態的,並非來自Storage,不需要被釋放 bool delay_alloc; //資料分配是否需要延緩,注意對輔助資料aux data無效 NDArrayStorageType storage_type = kDefaultStorage; std::vector<int> aux_types; Context ctx; TShape storage_shape; std::vector<TShape> aux_shapes; };
可見,Chunk結構仍然不是最終的資料儲存結構,本質上資料還是儲存在Storage結構中,如下所示:
graph LR NDArray-->|使用|Chunk Chunk-->|使用|Storage
在ndarray中,我們發現數據分為資料本身,以及輔助資料。輔助資料主要用於儲存稀疏資料的時候,資料本身放在data中,資料索引放在aux_data中。
最後看下NDArray的資料結構:
class NDArray { std::shared_ptr<Chunk> ptr_{nullptr}; TShape shape_; size_t byte_offset_ = 0; int dtype_ = -1; bool reuse_ = false; nnvm::NodeEntry entry_; mutable TBlob tblob_; };
resource
在mxnet中,計算中用到的所有內容,除了ndarray之外,都可以被稱為資源。其中最常用的資源,就是隨機數生成器,分為CPU和GPU兩個版本,如下:
enum Type { kRandom, //CPU版本隨機數生成器 kTempSpace, //動態隨機記憶體 kParallelRandom //可以在GPU中使用的並行隨機數生成器 };
另外,mxnet還為資源提供了一個管理器,ResourceManager,用於獲取資源。
kvstore
kv儲存的作用是儲存模型引數,以便在分散式的計算中,在多個裝置/機器之間進行資料同步。
kv儲存可以有多種型別,比如:
- 'local'或者'local_update_cpu‘或者'local_allreduce_cpu',表明這是一個單機的kv儲存,並且僅使用cpu做kv的allreduce;
- 'device'或者'local_allreduce_device',也是單機的kv儲存,只不過使用gpu做kv的allreduce;
- 'dist_*',分散式的kv儲存;
每個kv儲存中都有一個更新器,它定義了,針對指定的key,當新value來到時,如何與舊value進行融合。這一點非常重要,因為在深度學習模型的訓練中,需要迭代式的對模型引數進行更新,而更新的方式就是通過更新器來定義。
kv儲存中,key通常是整型或者字串,而value是NDArray,因此,有兩種更新器的定義:
typedef std::function<void(int, const NDArray&, NDArray*)> Updater; typedef std::function<void(const std::string&, const NDArray&, NDArray*)> StrUpdater;
最後,kv儲存在底層用到了ps-lite來作資料同步。
class KVStore { public: static KVStore *Create(const char *type = "local"); virtual void Init(const std::vector<int>& keys, const std::vector<NDArray>& values) = 0; virtual void Init(const std::vector<std::string>& str_keys, const std::vector<NDArray>& values) = 0; virtual void Push(...) = 0; virtual void Pull(...) = 0; virtual void PullRowSparse(...) = 0; virtual void set_updater(...); };
base
引入了兩個類,執行環境的上下文資訊類Context,實際執行時的上下文類RunContext,後者包含前者。首先看下Context類的定義:
struct Context { DeviceType dev_type; int32_t dev_id; inline void Save(dmlc::Stream *strm) const {...}; //將Context資訊記入二進位制流 inline bool Load(dmlc::Stream *strm) {...}; //從二進位制流中載入Context資訊 inline static Context Create(DeviceType dev_type, int32_t dev_id = -1); //構造一個新的Context inline static Context CPU(int32_t dev_id = 0); inline static Context GPU(int32_t dev_id=-1); inline static int32_t GetGPUCount(); //獲取GPU的數量 inline static void GetGPUMemoryInformation(int dev, int *free, int *total); inline static Context CPUPinned(int32_t dev_id = -1); inline static Context CPUShared(int32_t dev_id = 0); inline static Context FromString(const std::string& str); };
而RunContext就相對簡單了,它包含了一個Context和一個流指標:
struct RunContext { Context ctx; void *stream; //... };
operator
Operator定義了mxnet計算圖中基礎的操作單位。相當於Tensorflow中的kernel,和Caffe2中的Operator。但它與Tensorflow和Caffe2中的操作有本質區別,在Tensorflow中,操作本身和它對應的求導操作是分開的,而在mxnet中,這兩者是結合在一起的,分別使用Forward和Backward兩個函式實現,因此,mxnet在操作的實現上更加緊湊,與Tensorflow相比減少了一些對計算圖進行裁剪的額外開銷,效能上有優勢,但也同時限制了自己的計算邊界,靈活性不足。
class Operator { public: //進行前向計算,將計算結果儲存在TBlob中 virtual void Forward(const OpContext &ctx, const std::vector<TBlob> &in_data, const std::vector<OpReqType> &req, const std::vector<TBlob> &out_data, const std::vector<TBlob> &aux_states) = 0; //進行後向計算,將梯度寫入in_grad virtual void Backward(const OpContext &ctx, const std::vector<TBlob> &out_grad, const std::vector<TBlob> &in_data, const std::vector<TBlob> &out_data, const std::vector<OpReqType> &req, const std::vector<TBlob> &in_grad, const std::vector<TBlob> &aux_states); };
Operator中僅包含了操作計算的介面,對於操作的描述儲存在OperatorProperty類中,它負責儲存所有與Operator有關的資訊,且能夠產生裝置相關的Operator。同時,它還為計算引擎提供了一些可以優化操作計算的函式。
class OperatorProperty { public: //初始化Operator時需要用到的引數 virtual void Init(const std::vector<std::pair<std::string, std::string>>& kwargs) = 0; //獲取為Operator準備的引數 virtual std::map<std::string, std::string> GetParams() const = 0; virtual int NumOutputs() const {...} //進行Operator的形狀推斷,類似於Tensorflow的ShapeInference virtual bool InferShape(std::vector<TShape> *in_shape, std::vector<TShape> *out_shape, std::vector<TShape> *aux_shape) const = 0; //進行Operator的型別推斷 virtual bool InferType(...); //構建Operator virtual Operator* CreateOperator(Context ctx) const = 0; };
目前看來,mxnet中Operator與OperatorProperty的關係,與Tensorflow中OpKernel與Op的關係不太一樣,後者與Caffe2中的Operator和OpSchema的關係更加相似,有機會我們會詳細比較下,這三種框架關於操作定義於使用的區別。
engine
引擎是執行核心之一,它負責對計算圖中的操作進行排程。引擎中的兩大關鍵元素是操作和變數,操作定義了計算圖每一個節點需要實際執行的動作,變數定義了動作之間的依賴關係。
首先,mxnet定義了一個,被非同步函式在執行結束時呼叫的回撥函式類,通過對()的過載,用類對回撥函式進行了一層封裝:
class CallbackOnComplete { public: inline void operator()() const { (*callback_)(engine_, param_); } private: friend class ::mxnet::Engine; void (*callback_)(Engine *, void *); Engine* engine_; void* param_; };
列舉類FnProperty介紹了常用的函式型別:
enum class FnProperty { kNormal, //一般操作 kCopyFromGPU, //從GPU上拷貝內容到其它裝置的操作 kCopyToGPU, //從其它裝置向GPU拷貝內容的操作 kCPUPrioritized, //CPU上優先選擇的同步操作 kAsync, //非同步函式呼叫 kDeleteVar, //用來刪除變數的函式 kGPUPrioritized, //GPU上優先選擇的同步操作 };
engine的含義是,對操作進行排程執行的引擎。回想一下,在Tensorflow中,為了正確執行使用者設計好的計算圖,我們需要對原始計算圖進行一些迭代修改,在Engine類中提供了這樣的介面:
class Engine { public: //定義執行結束時的回撥類 typedef engine::CallbackOnComplete CallbackOnComplete; //定義傳遞給引擎的同步操作函式 typedef std::function<void(RunContext)> SyncFn; //定義傳遞給引擎的非同步操作函式 typedef std::function<void(RunContext, CallbackOnComplete)> AsyncFn; //定義變數指標 typedef engine::VarHandle VarHandle; //定義操作指標 typedef engine::OprHandle OprHandle; //停止引擎中的所有worker virtual void Stop() {} //啟動引擎中的所有worker virtual void Start() {} //分配一個新的變數,該變數可以被用來根據依賴關係,輔助對引擎中的操作進行排程 virtual VarHandle NewVariable() = 0; //構建一個操作,該操作定義在外部,從而我們可以在排程中重複使用 virtual OprHandle NewOperator(...) = 0; //刪除一個操作,它不會立刻進行,而是直到所有使用該操作的動作執行結束之後再進行 virtual void DeleteOperator(OpHandle op) = 0; //將一個操作加入引擎 virtual void Push(...); //將一個非同步操作加入引擎 virtual void PushAsync(...); //將一個同步操作加入引擎 virtual void PushSync(...); //刪除一個變數,它不會立刻進行,而是直到所有依賴該變數的操作完成之後再進行 virtual void DeleteVariable(...) = 0; //等待一個變數準備完成 virtual void WaitForVar(...) = 0; //等待引擎中所有的活動都結束時再返回 virtual void WaitForAll() = 0; //返回引擎的單例物件 static Engine* Get(); //用來生成OnComplete回撥的工廠函式 inline CallbackOnComplete CreateCallback(...); };
executor
mxnet的執行器介面,用於對計算圖進行執行。執行的機制與Operator的設計相合,同樣提供了前向和後向兩種介面,如下:
class Executor { public: virtual void Forward(bool is_train) = 0; virtual void PartialForward(bool is_train, int step, int *step_left) = 0; virtual void Backward(const std::vector<NDArray> &head_grads, bool is_train = true) = 0; };
rtc
包含了Cuda執行時的編譯模組CudaModule。
graph_attr_types
獲取圖相關屬性的輔助結構。對於一張計算圖中的節點,通常會關注兩種資訊,一種是計算圖中節點的儲存型別,一種是節點的排程模式,分別將結果儲存在StorageTypeVector和DispatchModeVector中,這兩種結構的定義如下:
using StorageTypeVector = std::vector<int>; using DispatchModeVector = std::vector<DispatchMode>;
op_attr_types
有關操作的額外屬性,與nvvm有關,目前看不懂。
imperative
與NDArray有關的執行時函式,目前看不懂。
operator_util
輔助快速構建operator的功能和註冊器。
c_api
定義了mxnet後端"C++"程式碼的介面。