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-->|包含|DLTensorndarray
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++"程式碼的介面。