Apache MXNet - 统一操作员 API

---

本章提供有关 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'是平滑的 l1 损失函数，我们假设它位于mshadow中。

看起来不可能将这种特定的损失实现为一元或二元运算符，但 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 <mxnet/operator_util.h>
   #if defined(__CUDACC__)
      #define XPU gpu
   #else
      #define XPU cpu
#endif

在我们的smooth l1 loss 示例中，输出与源具有相同的形状，我们可以使用默认Behave。它可以写成如下 -

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<gpu>* in GPU mode
   template<typename xpu> inline mshadow::Stream<xpu>* 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<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)));
   });
}

定义渐变

除了Input、TBlob和OpReqType加倍之外，二元运算符的梯度函数具有相似的结构。让我们看看下面，我们创建了一个具有各种类型输入的梯度函数：

// 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 损失示例。为了启用梯度链式法则，我们需要将顶部的out_grad乘以in_grad的结果。

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

将 SimpleOp 注册到 MXNet

创建形状、函数和梯度后，我们需要将它们恢复为 NDArray 运算符以及符号运算符。为此，我们可以使用注册宏，如下所示 -

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

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

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

EnvArguments 上的 SimpleOp

据我们所知，某些操作可能需要以下内容 -

• 标量作为输入，例如梯度比例

• 一组控制Behave的关键字参数

• 加速计算的临时空间。

使用 EnvArguments 的好处是它提供了额外的参数和资源，使计算更具可扩展性和效率。

例子

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

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

接下来，我们需要从EnvArguments.resource请求额外的资源，例如mshadow::Random<xpu>和临时内存空间。可以按如下方式完成 -

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 请求声明的资源请求。之后，它将资源放置在EnvAguments 中的 std::vector<Resource> 资源中。

我们可以借助以下代码访问资源 -

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 损失示例中看到，需要标量输入来标记损失函数的转折点。这就是为什么在注册过程中，我们在函数和梯度声明中使用set_enable_scalar(true)和env.scalar。

构建张量运算

这里出现了一个问题：为什么我们需要精心设计张量运算？原因如下：

• 计算使用 mshadow 库，但有时我们没有现成的函数。

• 如果操作不是以元素方式完成的，例如 softmax 损失和梯度。

例子

在这里，我们使用上面的 smooth 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;
         }
      }
   };
}