继续补充MXNet的操作符的内容
统一NDArray Opeartor和符号Opeartor
两者相似,区别在于后者有一个完整的依赖图;而逻辑底层基本是相同的;
SimpleOp 新的统一化运算符API
为I了构造一个依赖图,你必须搞懂:
是否有输出值、输入数据,或者头部梯度之外什么都不需要
统一操作API的梯度函数,会随着计算的操作类型的差异而辨别。
在了解更多的统一操作符API之前,首先学习下mshadow library guide
计算都为在mshasow::TBlob结构体中
实例:创建一个运算符函数smooth l1 loss,该函数混合了 l1 loss and l2 loss
loss = outside_weight .* f(inside_weight .* (data - label))
grad = outside_weight .* inside_weight .* f'(inside_weight .* (data - label))
.*表示:元素相乘
SimpleOP:统一的运算符API
定义Shapes
mshadow库需要显示的内存分配,因而,所有的DataShapes必须在计算发生前准备完毕;
在我们用定义好的函数和梯度处理前,先检查输入的Shape的一致性以及输出的shape
typedef TShape (*UnaryShapeFunction)(const TShape& src, const EnvArguments& env);
typedef TShape (*BinaryShapeFunction)(const TShape& const TShape& rhs,lhs,
const EnvArguments& env);
备注:感受此处定义有误,待后续矫正。。
利用mshadow::TShape来检查input Data Shape,并且指派输出Data Shape.
如果不定义该函数,则Input Shape和Output Shape类型一致;
也可以用shape函数来检查附加的参数和资源是可用的;同时关注EnvArguments
#include <mxnet/operator_util.h>
#if defined(__CUDACC__)
#define XPU gpu
#else
#define XPU cpu
#endif
在开始smooth l1 loss实例前,定义一个XPU to cpu or cpu在smooth_l1_unary-inl.h实例中
这样来实现重用 smooth_l1_unary.cc
和smooth_l1_unary.cu的相同代码
在smooth l1 loss实例中,可以使用缺省的行为;
inline TShape SmoothL1Shape_(const TShape& src,
const EnvArguments& env) {
return TShape(src);
定义函数
通过mshadow:TBlob来创建
一个unary或者binary函数
typedef void (*UnaryFunction)(const TBlob& src,
const EnvArguments& env,
TBlob* ret,
OpReqType req,
RunContext ctx);
typedef void (*BinaryFunction)(const TBlob& lhs,
const TBlob& rhs,
const EnvArguments& env,
TBlob* ret,
OpReqType req,
RunContext ctx);
函数会根据输入参数的类型,而具备差异性
关注下RunContext,在执行器的运行时是必须的
struct RunContext {
void *stream; // the stream of the device, can be NULL or Stream<gpu>* in GPU mode
template<typename xpu> inline mshadow::Stream<xpu>* get_stream() // get mshadow stream from Context
} // namespace mxnet
从ctx上获取stream的实例
mshadow::stream<xpu> *s = ctx.get_stream<xpu>();
OpReqType req 表示计算结果如何写入到ret中
enum OpReqType {
kNullOp, // no operation, do not write anything
kWriteTo, // write gradient to provided space
kWriteInplace, // perform an in-place write
kAddTo // add to the provided space
};
在operator_util.h中定义了宏,来简化OpReqType.ASSIGN_DISPATCH(out,req,exp)
并且检查req来执行赋值
template<typename xpu>
void SmoothL1Forward_(const TBlob& src,
const EnvArguments& env,
TBlob *ret,
OpReqType req,
RunContext ctx) {
using namespace mshadow;
using namespace mshadow::expr;
mshadow::Stream<xpu> *s = ctx.get_stream<xpu>();
real_t sigma2 = env.scalar * env.scalar;
MSHADOW_TYPE_SWITCH(ret->type_flag_, DType, {
mshadow::Tensor<xpu, 2, DType> out = ret->get<xpu, 2, DType>(s);
mshadow::Tensor<xpu, 2, DType> in = src.get<xpu, 2, DType>(s);
ASSIGN_DISPATCH(out, req,
F<mshadow_op::smooth_l1_loss>(in, ScalarExp<DType>(sigma2)));
});
}
注意:
MSHADOW_TYPE_SWITCH(type,DType,...)处理不同类型的详细
ASSIGN_DISPATCH(out,req,exp)检查OpReqType,并执行相关action
定义梯度
可以根据输入的多种类型,来创建一个gradient函数
// depending only on out_grad
typedef void (*UnaryGradFunctionT0)(const OutputGrad& out_grad,
const EnvArguments& env,
TBlob* in_grad,
OpReqType req,
RunContext ctx);
// depending only on out_value
typedef void (*UnaryGradFunctionT1)(const OutputGrad& out_grad,
const OutputValue& out_value,
const EnvArguments& env,
TBlob* in_grad,
OpReqType req,
RunContext ctx);
// depending only on in_data
typedef void (*UnaryGradFunctionT2)(const OutputGrad& out_grad,
const Input0& in_data0,
const EnvArguments& env,
TBlob* in_grad,
OpReqType req,
RunContext ctx);
二进制运算符的梯度函数具有类型的结构;处理Input,TBlob,OpReqType是doubled
GradFunctionArgument,Input0,Input,OutputValue和OutputGrad都共享GradFunctionArgument结构
struct GradFunctionArgument {
TBlob data;
}
关注smooth l1 loss实例
template<typename xpu>
void SmoothL1BackwardUseIn_(const OutputGrad& out_grad,
const Input0& in_data0,
const EnvArguments& env,
TBlob *in_grad,
OpReqType req,
RunContext ctx) {
using namespace mshadow;
using namespace mshadow::expr;
mshadow::Stream<xpu> *s = ctx.get_stream<xpu>();
real_t sigma2 = env.scalar * env.scalar;
MSHADOW_TYPE_SWITCH(in_grad->type_flag_, DType, {
mshadow::Tensor<xpu, 2, DType> src = in_data0.data.get<xpu, 2, DType>(s);
mshadow::Tensor<xpu, 2, DType> ograd = out_grad.data.get<xpu, 2, DType>(s);
mshadow::Tensor<xpu, 2, DType> igrad = in_grad->get<xpu, 2, DType>(s);
ASSIGN_DISPATCH(igrad, req,
ograd * F<mshadow_op::smooth_l1_gradient>(src, ScalarExp<DType>(sigma2)));
});
}
注意是一个f'(x),利用输入来进行梯度计算,UnaryGradFunctionT2是合适的;
为了保证梯度的链条规则,我们需要从in_grad的结果的头部来乘out_grad
把SimpleOP注册到MXNet
在定义完shape,function,gradient后,可以存储为NDArray或者Symbolic操作符;为了简化这步处理,可以用operator_util.h的宏
MXNET_REGISTER_SIMPLE_OP(Name, DEV)
.set_shape_function(Shape)
.set_function(DEV::kDevMask, Function<XPU>, SimpleOpInplaceOption)
.set_gradient(DEV::kDevMask, Gradient<XPU>, SimpleOpInplaceOption)
.describe("description");
SimpleOpInplaceOption定义如下:
enum SimpleOpInplaceOption {
kNoInplace, // do not allow inplace in arguments
kInplaceInOut, // allow inplace in with out (unary)
kInplaceOutIn, // allow inplace out_grad with in_grad (unary)
kInplaceLhsOut, // allow inplace left operand with out (binary)
kInplaceOutLhs // allow inplace out_grad with lhs_grad (binary)
};
关注inplace
我们有一个梯度函数,依赖输入数据;所有该function不能in place写入;输出梯度在梯度计算后没有其它用处;可以in place写入。
MXNET_REGISTER_SIMPLE_OP(smooth_l1, XPU)
.set_function(XPU::kDevMask, SmoothL1Forward_<XPU>, kNoInplace)
.set_gradient(XPU::kDevMask, SmoothL1BackwardUseIn_<XPU>, kInplaceOutIn)
.set_enable_scalar(true)
.describe("Calculate Smooth L1 Loss(lhs, scalar)");
在shape函数的讨论中,没有set_shape_function的缺省行为,要求输入具有相同的shape,然后再生成相同的shape作为输出;
我们会后续讨论set_enable_scalar。
NDArray运算符汇总
1.为了决定输出shape来创建一个shape函数
2.通过选择一个合适的函数类型,创建一个函数作为forward 例行程序
3.通过选择一个合适的梯度类型,创建一个梯度作为backward例行程序
4.通过注册process来注册运算符
SimpleOP的附件信息
在环境参数上使用SimpleOP
一些运算符需要一个标量作为输入,比如梯度规模,一套控制行为的关键词参数,一个临时空间来加速计算
可以通过EvnArguments提供额外的参数和资源,来进行可伸缩性和高效的计算。
struct EnvArguments {
real_t scalar; // scalar argument, if enabled
std::vector<std::pair<std::string, std::string> > kwargs; // keyword arguments
std::vector<Resource> resource; // pointer to the resources requested
};
为了这些附加的特性,需要更多的注册参数;
为了避免参数的混乱性,scalar和kwargs不能够同时出现
为了使scalar生效,使用set_enable_scalar(bool enable_scalar)来注册,然后在前向的梯度函数中,scalar能够从env.scala获取,通过EnvArguments env
为了使kwargs生效,使用set_enable_kwargs(bool enable_kwargs)来注册,然后在前向的梯度函数中,额外的参数包含在env.kwarg中。
而这些函数定义在std::vector<std::pair<std::string, std::string> >中,通过DMLC参数构造来简化关键字参数的解析
关注parameter structure
额外的资源像mshadow::Randow<xpu> 和临时的内存空间可以通过EnvArguments.resource来存取和获得请求;
注册程序是:
set_resource_request(ResourceRequest req)或者set_resource_request(const std::vector<ResourceRequest>)
定义如下:
struct ResourceRequest {
enum Type { // Resource type, indicating what the pointer type is
kRandom, // mshadow::Random<xpu> object
kTempSpace // A dynamic temp space that can be arbitrary size
};
Type type; // type of resources
};
注册请求是需要被声明的资源请求的,从mxnet::ResourceManager
并且把资源安置在std::vector<Resource> resource和EnvArguments
为了获取资源,通过
auto tmp_space_res = env.resources[0].get_space(some_shape, some_stream);
auto rand_res = env.resources[0].get_random(some_stream);
为此:
在smooth l1 loss实例中,标识loss函数的转换点必须一个 scalar input
所以,在注册的处理过程中,我们使用set_enable_scalar(true),并且在梯度和函数声明中使用 env.scalar 。
手工制作一个Tensor运算符
计算利用mshadow库,所以我们能做operator实现中来创建一个tensor运算符。
如果想定义逐个元素的函数. 可以用mxnet::op::mshadow_op /src/operator/mshadow_op.h来实现。
可以参考mshadow experssion API guide
如果一个操作符,不能通过逐个元素的方式定义;想softmax loss和gradient,需要创建一个新的tensor运算符;
比如mshadow::cuda,创建两个mappers
namespace mshadow_op {
struct smooth_l1_loss {
// a is x, b is sigma2
MSHADOW_XINLINE static real_t Map(real_t a, real_t b) {
if (a > 1.0f / b) {
return a - 0.5f / b;
} else if (a < -1.0f / b) {
return -a - 0.5f / b;
} else {
return 0.5f * a * a * b;
}
}
};
}
超出2个运算数
新的统一化API被用来设计作为实现运算符的基础;针对两个输出参数以上的运算符;可以参考Opeartor API