Apache MXNet - 统一算子 API

上一页
下一页

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

SimpleOp

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 <mxnet/operator_util.h>
   #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<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)");

SimpleOp 上的 EnvArguments

众所周知，某些操作可能需要以下内容：

  • 一个标量作为输入，例如梯度缩放

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

  • 一个临时空间来加速计算。

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

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

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;
         }
      }
   };
}
打印页面
上一页
下一页
广告

© . All rights reserved.