您现在的位置是:网站首页> 内容页

mxnet源码阅读笔记之include

  • 登录沙龙3salon366
  • 2019-05-29
  • 266人已阅读
简介写在前面mxnet代码的规范性比Caffe2要好,看起来核心代码量也小很多,但由于对dmlc其它库的依赖太强,代码的独立性并不好。依赖的第三方库包括:cubdlpackdmlc-cor

写在前面

mxnet代码的规范性比Caffe2要好,看起来核心代码量也小很多,但由于对dmlc其它库的依赖太强,代码的独立性并不好。依赖的第三方库包括:

cubdlpackdmlc-coregoogletestmkldnnmshadowonnx-tensorrtopenmpps-litetvm

如果对于这些第三方库没有足够的理解,mxnet的核心代码看起来比较费劲。因此时间原因,本篇仅解析了mxnet对外的接口include目录,并且对于严重依赖第三方库的文件没有深入探究,只能算作一篇不完整的源码阅读笔记了。后续有时间的话,再回来迭代。

目录

storagetensor_blobndarrayresourcekvstorebaseoperatorengineexecutorrtcgraph_attr_typesop_attr_typesimperativeoperator_utilc_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++"代码的接口。

, 1, 0, 9);

文章评论

Top