在MXNet中,运算符是一个包含实际计算逻辑和辅助信息的类,可以帮助系统执行优化,如就地更新和自动求导。在继续使用本文档之前,我们强烈建议您了解mshadow
库,因为所有运算符都是在运行时由系统提供的张量结构mshadow :: TBlob
计算的。
MXNet的运算符接口允许您:
cudnn
例程)。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;
“OpContext”结构是:
struct OpContext {
int is_train
RunContext run_ctx;
std :: vector <Resource> requested;
}}
它描述了运算符是处于训练还是测试阶段(在is_train
中指定),运算符应该在哪个设备上运行(在run_ctx
中),以及请求的资源(在以下部分中介绍)。
in_data
和out_data
分别表示输入和输出张量。所有的张量空间都由系统分配。
req
表示如何将计算结果写入out_data
。换句话说,req.size()== out_data.size()
和req [i]
对应于out_data [i]
的写类型。
运算符请求类型OpReqType定义为:
enum OpReqType {
kNullOp,
kWriteTo,
kWriteInplace,
kAddTo
};
通常,所有out_data
的类型应该是kWriteTo
,意味着提供的out_data
张量是一个原始内存块,因此运算符应该将结果直接写入它。在某些情况下,例如当计算梯度张量时,如果我们可以累积结果,而不是直接重写张量内容,这样就很好,这样每次不需要创建额外的空间。在这种情况下,相应的req
类型被设置为kAddTo
,表示应该调用+ =
。当我们需要进行同址运算(in-place computation)时,out_data
的类型应该是kWriteInplace
,这时输出张量out_data
将与输入张量in_data
共享同一片内存,将运算后的结果就地写入输入张量,而不是新开辟一块内存,以节省内存开销。
aux_states
被有意设计用于帮助计算的辅助张量。目前,它是无用的。除了Forward
运算符,你可以选择实现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_data”和“out_data”,并且运算符计算“in_grad”作为结果。命名策略类似于Torch的约定,可以总结如下图(mxnet官方没有给出图片):
[输入/输出语义图]
一些运算符可能不需要所有以下内容:out_grad
,in_data
和out_data
。这些可以在OperatorProperty
中的DeclareBackwardDependency
接口中指定。
一个卷积可能有几个实现,你可能想在它们之间切换以实现最佳性能。因此,我们将operator 的* 语义 *接口从实现接口(Operator
类)剥离为OperatorProperty
类。 OperatorProperty
接口包括:
virtual bool InferShape(std :: vector <TShape> * in_shape,
std :: vector <TShape> * out_shape,
std :: vector <TShape> * aux_shape)const = 0;
这个接口有两个目的:1.告诉系统每个输入和输出张量的大小,所以它可以在Forward
和Backward
调用之前为它们分配空间; 2.执行大小检查以确保在运行之前没有明显的错误。 in_shape
中的形状将由系统设置(根据前面的运算符的out_shape
)。当没有足够的信息来推断形状时,它返回false
,当形状不一致时抛出错误。
cudnnConvolutionForward
这样的运算符需要一个用于计算的工作空间。如果系统可以管理它,它可以执行优化,如重用该空间,等等。 MXNet定义了两个接口来实现: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结构(在resource.h中)目前只包含一个类型标志:
struct ResourceRequest {
enum Type{
kRandom,// 获得一个mshadow :: Random <xpu> 对象
kTempSpace,//请求临时空间
};
Type type;
};
如果ForwardResource
和BackwardResource
返回非空数组,系统通过Operator
的Forward
和Backward
接口中的ctx
参数提供相应的资源。基本上,要访问这些资源,只需写:
auto tmp_space_res = ctx.requested [kTempSpace] .get_space(some_shape,some_stream);
auto rand_res = ctx.requested [kRandom] .get_random(some_stream);
有关示例,请参见src / operator / cudnn_convolution-inl.h
。
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);
注意,FullyConnectedForward
中的out_data
不被FullyConnectedBackward
使用,PoolingBackward
需要PoolingForward
的所有参数。因此,对于FullyConnectedForward
,out_data
张量一旦消耗可以安全地释放,因为向后的函数不需要它。这提供了系统尽可能快地对一些张量进行垃圾回收的机会。为了指定这种情况,我们提供了一个接口:
virtual std :: vector <int> DeclareBackwardDependency(
const std :: vector <int>&out_grad,
const std :: vector <int>&in_data,
const std :: vector <int>&out_data)const;
参数向量的int
元素是用于区分不同数组的ID。让我们看看这个接口如何为FullyConnected
和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]}; //注:不包括out_data [0]
}}
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]};
}}
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]} }
}
这告诉系统在Forward
期间in_data [0]
和out_data [0]
张量可以共享相同的存储空间,out_grad [0]
和in_grad [0]
在Backward
期间共享内存。
重要:即使您使用上述规范,也不能保证输入和输出张量共享相同的空间。事实上,这只是对系统作最后的决定的一个建议。但是,在任何一种情况下,这个决定对你是完全透明的,所以实际的“Forward”和“Backward”实现不需要考虑。
//从键值字符串列表中初始化属性类
virtual void Init(const vector <pair <string,string >>&kwargs)= 0;
//返回键值字符串映射中的参数
virtual map <string,string> GetParams()const = 0;
//返回参数的名称(用于在python中生成签名)
virtual vector<string> ListArguments()const;
//返回输出值的名称
virtual vector<string> ListOutputs()const;
//返回辅助状态的名称
virtual vector<string> ListAuxiliaryStates()const;
//返回输出值的个数
virtual int NumOutputs()const;
//返回可见输出的数量
virtual int NumVisibleOutputs()const;
OperatorProperty
包括操作的所有语义属性。它还负责创建用于实际计算的“运算符”指针。
在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;
}}
};
当实现卷积运算符时,您需要知道卷积核大小,步幅大小,填充大小等。在调用任何Forward
或Backward
接口之前,这些参数应该传递给操作符。为此,您可以定义一个ConvolutionParam
结构,如下所示:
#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_;
};
使用以下宏来将参数结构和运算符属性类注册到MXNet:
DMLC_REGISTER_PARAMETER(ConvolutionParam);
MXNET_REGISTER_OP_PROPERTY(Convolution,ConvolutionOpProperty);
第一个参数是名称字符串,第二个是属性类名。
我们几乎覆盖了定义新运算符所需的整个接口。让我们回顾一下:
Operator
接口来写你的计算逻辑(Forward
和Backward
)。OperatorProperty
接口:- 将参数传递给操作符类(可以使用Init
接口)。
- 使用CreateOperator
接口创建一个操作符。
- 正确实现操作符描述接口,例如参数名等。
- 正确实现“InferShape”接口以设置输出张量形状。
- [可选]如果需要额外的资源,请选中“ForwardResource”和“BackwardResource”。
- [可选]如果Backward
不需要Forward
的所有输入和输出,请检查DeclareBackwardDependency
。
- [可选]如果支持就地更新,请检查“ForwardInplaceOption”和“BackwardInplaceOption”。
OperatorProperty
类和参数类。