一个运算符Operator是一个类,同时包含计算逻辑和辅助信息(帮助执行性能优化,比如及时更新和自动求导)
为此,建议参考mshadow库,类似tensor结构的mshadown:TBlob
MXNet运算符接口的用处:
通过指定及时更新来减少内存分配
从Python来隐藏内部参数
定义输入和输出tensors的关系,允许系统执行模型检查
从系统中获得附加的临时空间执行计算;比如调用cudnn routines
运算符接口(Operator Inteface)
Forward 核心运算符号
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;
上下文
struct OpContext { int is_train; RunContext run_ctx; std::vector<Resource> requested; }
定义:运算符所在的阶段(Train,Test)
运算符在什么设备上执行,以及所需的资源。
req指示:计算结果如何写入out_data
req.size() == out_data.size()
并且req[i]对应到out_data[i]
and req[i] correspond to the write type of out_data[i]
运算请求类型OpReqType
enum OpReqType { kNullOp, kWriteTo, kWriteInplace, kAddTo };
kWriteTo和kAddTo之间的差异化
通常情况下,所有的out_data都是kWriteTo
但是针对计算梯度下降的Tensor,计算量非常巨大,通常采用递增的方式处理
aux_states用来设计支持辅助的Tensor的,目前用处不大。
Backward接口
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);
该接口类似Forward,除了out_grad,in_grad
注意:in_grad是结果
命名策略类似Torch的规约 [input/output semantics figure]
针对部分运算符,不需要out_grad,in_data,out_data;可以用DeclareBackwardDependency来指定依赖关系。
Operator Property[运算符的属性]
卷积有许多实现,你或者想在这些卷积之间进行网络切换,并获取最好的性能;
为此,分离运算符的语义接口和实现接口;Opeartor 和 OpeartorProperty
关注OperatorProperty的核心
InferShape 推导模型
virtual bool InferShape(std::vector<TShape> *in_shape, std::vector<TShape> *out_shape, std::vector<TShape> *aux_shape) const = 0;
该接口的主要用途
1.告诉每个Input和Output的Tensor的大小,在执行Forward和Backward之前分配好空间。
2.执行大小检查,in_shape的由系统来设置,根据之前的out_shape来完成;如果没有足够多的信息来推导的话,会返回false
如果模型不一致,则抛出异常。
Request Resources
类似cudnnConvolutionForward运算符需要专门的计算空间;如果一个系统能够管理,就可以执行优化,进行重复利用。
virtual std::vector<ResourceRequest> ForwardResource( const std::vector<TShape> &in_shape) const;virtual std::vector<ResourceRequest> BackwardResource( const std::vector<TShape> &in_shape) const;
注意结构体 ResourceRequest 的定义
struct ResourceRequest { enum Type { kRandom, // get a mshadow::Random<xpu> object kTempSpace, // request temporary space }; Type type; };
如果前向和后向资源返回非空数组;系统会通过接口中的ctx参数来提供相关资源;
资源的获取方式:
auto tmp_space_res = ctx.requested[kTempSpace].get_space(some_shape, some_stream);auto rand_res = ctx.requested[kRandom].get_random(some_stream);
反向依赖Backward dependency
关注下两类不同类型的操作符定义
void FullyConnectedForward(TBlob weight, TBlob in_data, TBlob out_data);void FullyConnectedBackward(TBlob weight, TBlob in_data, TBlob out_grad, TBlob in_grad);void PoolingForward(TBlob in_data, TBlob out_data);void PoolingBackward(TBlob in_data, TBlob out_data, TBlob out_grad, TBlob in_grad);
注意FullyConnectedFoward的out_data没有在FullyConnectedBackward中使用
而PoolingBackward会全部使用PoolingForward的参数
结果:FullyConnectedForward的out_data在一次消费后,就可以进行释放了;可以通过GC来完成
virtual std::vector<int> DeclareBackwardDependency( const std::vector<int> &out_grad, const std::vector<int> &in_data, const std::vector<int> &out_data) const;
参数Vector中的int元素,是一个区分不同数组的ID
关注这些接口如何指定FullyConnect和Pooling之间的不同依赖
std::vector<int> FullyConnectedProperty::DeclareBackwardDependency( const std::vector<int> &out_grad, const std::vector<int> &in_data, const std::vector<int> &out_data) const { return {out_grad[0], in_data[0]}; // NOTE: out_data[0] is NOT included }
std::vector<int> PoolingProperty::DeclareBackwardDependency( const std::vector<int> &out_grad, const std::vector<int> &in_data, const std::vector<int> &out_data) const { return {out_grad[0], in_data[0], out_data[0]}; }
In Place Option
为了保持内存分配的成本,可以通过In Place Update.
当input和output向量具有相同的模型时,对于element-wise的运算符是恰当的;
virtual std::vector<std::pair<int, void*>> ElewiseOpProperty::ForwardInplaceOption( const std::vector<int> &in_data, const std::vector<void*> &out_data) const { return { {in_data[0], out_data[0]} }; }virtual std::vector<std::pair<int, void*>> ElewiseOpProperty::BackwardInplaceOption( const std::vector<int> &out_grad, const std::vector<int> &in_data, const std::vector<int> &out_data, const std::vector<void*> &in_grad) const { return { {out_grad[0], in_grad[0]} } }
针对系统来说,in_data[0]和out_data[0]在执行Forward时可以共享相同的内存;同理,在Backward时,out_grad[0]和in_grad[0]同样允许。
核心关键:针对依赖的执行,在系统层面是否能够共享内存,针对用户态需要做到透明化。
向Python暴露运算符
// initial the property class from a list of key-value string pairsvirtual void Init(const vector<pair<string, string>> &kwargs) = 0;// return the parameters in a key-value string mapvirtual map<string, string> GetParams() const = 0;// return the name of arguments (for generating signature in python)virtual vector<string> ListArguments() const;// return the name of output valuesvirtual vector<string> ListOutputs() const;// return the name of auxiliary statesvirtual vector<string> ListAuxiliaryStates() const;// return the number of output valuesvirtual int NumOutputs() const;// return the number of visible outputsvirtual int NumVisibleOutputs() const;
从操作属性中创建一个操作符
OperatorProperty包括一个Operation的所有语法属性,可以针对执行任务时,创建Operator指针。
Create Operator
通过实现OperatorProperty的接口
virtual Operator* CreateOperator(Context ctx) const = 0;
实例
class ConvolutionOp { public: void Forward( ... ) { ... } void Backward( ... ) { ... } }; class ConvolutionOpProperty : public OperatorProperty { public: Operator* CreateOperator(Context ctx) const { return new ConvolutionOp; } };
参数化运算符
在实现一个卷积运算符时,需要知道kernel size, stride size, padding size等配置参数;
并且在forward和backward之前把这些参数传给Opeartor;可以定义一个卷积参数
#include <dmlc/parameter.h> struct ConvolutionParam : public dmlc::Parameter<ConvolutionParam> { TShape kernel, stride, pad; uint32_t num_filter, num_group, workspace; bool no_bias; };
把该参数设置在ConvolutionOpProperty中,并且在构造器中传入
class ConvolutionOp { public: ConvolutionOp(ConvolutionParam p): param_(p) {} void Forward( ... ) { ... } void Backward( ... ) { ... } private: ConvolutionParam param_; }; class ConvolutionOpProperty : public OperatorProperty { public: void Init(const vector<pair<string, string>& kwargs) { // initialize param_ using kwargs } Operator* CreateOperator(Context ctx) const { return new ConvolutionOp(param_); } private: ConvolutionParam param_; };
把OperatorProperty Class 和Parameter类注册到MXNet
可以采用如下宏来完成
DMLC_REGISTER_PARAMETER(ConvolutionParam);MXNET_REGISTER_OP_PROPERTY(Convolution, ConvolutionOpProperty);
接口汇总
1.利用Opeartor接口来写forward,backward的计算逻辑
2.使用OpeartorProperty可以满足如下场景的使用
2.1利用Init接口来传递参数
2.2用CreateOpeartor接口来创建操作符
3.注册OpeartorProperty类和参数类