📌  相关文章
📜  Apache MXNet-统一操作员API

📅  最后修改于: 2020-12-10 04:46:55             🧑  作者: Mango

本章提供有关Apache MXNet中的统一运算符应用程序编程接口(API)的信息。

简单操作

SimpleOp是新的统一运算符API，它统一了不同的调用过程。一旦被调用，它将返回到运算符的基本元素。统一运算符专为一元和二进制运算而设计。这是因为大多数数学运算符涉及一个或两个操作数，而更多的操作数使与相关性相关的优化非常有用。

我们将借助示例来了解其SimpleOp统一运算符。在此示例中，我们将创建一个运算符作用为平滑的l1损失，它是l1和l2损失的混合。我们可以定义并写出如下损失：

loss = outside_weight .* f(inside_weight .* (data - label))
grad = outside_weight .* inside_weight .* f'(inside_weight .* (data - label))

在上面的例子中

  • 。*代表逐元素乘法

  • f，f’是我们假设在mshadow中的平滑l1损失函数。

看起来不可能以一元或二进制运算符实现这种特殊的损失，但是MXNet在符号执行中为其用户提供了自动区分，从而直接简化了f和f’的损失。这就是为什么我们可以将这种特殊损失作为一元运算运算符来实现的原因。

定义形状

众所周知，MXNet的mshadow库需要显式的内存分配，因此我们需要在进行任何计算之前提供所有数据形状。在定义函数和渐变之前，我们需要提供输入形状一致性和输出形状，如下所示：

typedef mxnet::TShape (*UnaryShapeFunction)(const mxnet::TShape& src,
const EnvArguments& env);
   typedef mxnet::TShape (*BinaryShapeFunction)(const mxnet::TShape& lhs,
const mxnet::TShape& rhs,
const EnvArguments& env);

函数mxnet :: Tshape用于检查输入数据形状和指定的输出数据形状。如果没有定义此函数，则默认输出形状将与输入形状相同。例如，在二进制运算符的情况下，默认情况下会检查lhs和rhs的形状是否相同。

现在让我们继续我们平滑的l1损失示例。为此，我们需要在标头实现smooth_l1_unary-inl.h中将XPU定义为cpu或gpu 。原因是要在smooth_l1_unary.cc和smooth_l1_unary.cu中重用相同的代码。

#include 
   #if defined(__CUDACC__)
      #define XPU gpu
   #else
      #define XPU cpu
#endif

就像我们的平滑l1损失示例一样，输出的形状与源的形状相同，我们可以使用默认行为。它可以写成如下-

inline mxnet::TShape SmoothL1Shape_(const mxnet::TShape& src,const EnvArguments& env) {
   return mxnet::TShape(src);
}

定义功能

我们可以使用一个输入创建一元或二进制函数，如下所示：

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 ctx结构，其中包含运行时执行所需的信息-

struct RunContext {
   void *stream; // the stream of the device, can be NULL or Stream* in GPU mode
   template inline mshadow::Stream* get_stream() // get mshadow stream from Context
} // namespace mxnet

现在，让我们看看如何将计算结果写入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
};

现在，让我们继续进行平滑的l1损失示例。为此，我们将使用UnaryFunction定义此运算符的函数，如下所示：

template
void SmoothL1Forward_(const TBlob& src,
   const EnvArguments& env,
   TBlob *ret,
   OpReqType req,
RunContext ctx) {
   using namespace mshadow;
   using namespace mshadow::expr;
   mshadow::Stream *s = ctx.get_stream();
   real_t sigma2 = env.scalar * env.scalar;
   MSHADOW_TYPE_SWITCH(ret->type_flag_, DType, {
      mshadow::Tensor out = ret->get(s);
      mshadow::Tensor in = src.get(s);
      ASSIGN_DISPATCH(out, req,
      F<:smooth_l1_loss>(in, ScalarExp(sigma2)));
   });
}

定义渐变

除了将Input，TBlob和OpReqType加倍外，二进制运算符的Gradients函数具有相似的结构。让我们在下面查看，我们在其中创建了具有各种输入类型的渐变函数：

// 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);

如上定义， Input0，Input，OutputValue和OutputGrad都共享GradientFunctionArgument的结构。它的定义如下-

struct GradFunctionArgument {
   TBlob data;
}

现在让我们继续进行平滑的l1损失示例。对于此，使梯度的链式法则，我们需要从顶部向in_grad的结果乘以out_grad。

template
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 *s = ctx.get_stream();
   real_t sigma2 = env.scalar * env.scalar;
      MSHADOW_TYPE_SWITCH(in_grad->type_flag_, DType, {
      mshadow::Tensor src = in_data0.data.get(s);
      mshadow::Tensor ograd = out_grad.data.get(s);
      mshadow::Tensor igrad = in_grad->get(s);
      ASSIGN_DISPATCH(igrad, req,
      ograd * F<:smooth_l1_gradient>(src, ScalarExp(sigma2)));
   });
}

将SimpleOp注册到MXNet

创建形状，函数和渐变后，我们需要将它们还原到NDArray运算符和符号运算符。为此，我们可以使用注册宏，如下所示：

MXNET_REGISTER_SIMPLE_OP(Name, DEV)
   .set_shape_function(Shape)
   .set_function(DEV::kDevMask, Function, SimpleOpInplaceOption)
   .set_gradient(DEV::kDevMask, Gradient, 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)
};

现在让我们继续进行平滑的l1损失示例。为此，我们有一个依赖于输入数据的渐变函数，因此该函数无法就地编写。

MXNET_REGISTER_SIMPLE_OP(smooth_l1, XPU)
.set_function(XPU::kDevMask, SmoothL1Forward_, kNoInplace)
.set_gradient(XPU::kDevMask, SmoothL1BackwardUseIn_, kInplaceOutIn)
.set_enable_scalar(true)
.describe("Calculate Smooth L1 Loss(lhs, scalar)");

EnvArguments上的SimpleOp

我们知道一些操作可能需要以下内容-

  • 标量作为输入，例如渐变标度

  • 一组控制行为的关键字参数

  • 加快计算速度的临时空间。

使用EnvArguments的好处是它提供了其他参数和资源，以使计算更具可伸缩性和效率。

例

首先让我们定义结构如下：

struct EnvArguments {
   real_t scalar; // scalar argument, if enabled
   std::vector<:pair std::string=""> > kwargs; // keyword arguments
   std::vector resource; // pointer to the resources requested
};

接下来，我们需要从EnvArguments.resource请求其他资源，例如mshadow :: Random 和临时内存空间。它可以做到如下-

struct ResourceRequest {
   enum Type { // Resource type, indicating what the pointer type is
      kRandom, // mshadow::Random object
      kTempSpace // A dynamic temp space that can be arbitrary size
   };
   Type type; // type of resources
};

现在，注册将从mxnet :: ResourceManager请求声明的资源请求。之后，它将把资源放在EnvAgruments的std :: vector 资源中。

我们可以在以下代码的帮助下访问资源-

auto tmp_space_res = env.resources[0].get_space(some_shape, some_stream);
auto rand_res = env.resources[0].get_random(some_stream);

如果您在平滑的l1损失示例中看到，则需要标量输入来标记损失函数的转折点。这就是为什么在注册过程中，我们在函数和渐变声明中使用set_enable_scalar(true)和env.scalar的原因。

建筑张量操作

这里出现的问题是，为什么我们需要进行张量运算?原因如下-

  • 计算利用了mshadow库，有时我们没有可用的函数。

  • 如果没有以元素方式(例如softmax损失和梯度)进行操作。

例

在这里，我们使用上面的平滑l1损失示例。我们将创建两个映射器，即平滑l1损失和梯度的标量情况：

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;
         }
      }
   };
}