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' 是平滑的 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 的形状是否相同。

现在让我们继续讨论 smooth 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 示例中一样，输出具有与源相同的形状，我们可以使用默认行为。它可以写成如下形式 −

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; // 设备的流，在 GPU 模式下可以为 NULL 或 Stream<gpu>*
    template<typename xpu> inline mshadow::Stream<xpu>* get_stream() // 从 Context 获取 mshadow 流
} // namespace mxnet

现在我们来看看如何将计算结果写入ret中。

enum OpReqType {
    kNullOp, // 无操作，不写入任何内容
    kWriteTo, // 将渐变写入提供的空间
    kWriteInplace, // 执行就地写入
    kAddTo // 添加到提供的空间
};

现在，让我们继续讨论smooth l1 loss 示例。为此，我们将使用 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 加倍外，二元运算符的梯度函数具有类似的结构。让我们在下面查看，我们创建了一个具有各种输入类型的梯度函数：

// 仅取决于 out_grad
typedef void (*UnaryGradFunctionT0)(const OutputGrad& out_grad,
   const EnvArguments& env,
   TBlob* in_grad,
   OpReqType req,
   RunContext ctx);
// 仅取决于 out_value
typedef void (*UnaryGradFunctionT1)(const OutputGrad& out_grad,
   const OutputValue& out_value,
   const EnvArguments& env,
   TBlob* in_grad,
   OpReqType req,
   RunContext ctx);
// 仅取决于 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;
}

现在让我们继续讨论smooth l1 loss 示例。为了启用梯度链式法则，我们需要将顶部的 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, // 不允许 inplace in 参数
    kInplaceInOut, // 允许 inplace in 和 out（一元）
    kInplaceOutIn, // 允许 inplace out_grad 和 in_grad（一元）
    kInplaceLhsOut, // 允许 inplace left 操作数和 out（二元）
    
    kInplaceOutLhs // 允许 inplace out_grad 和 lhs_grad（二元）
};

现在让我们继续讨论 平滑 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("计算 Smooth L1 Loss(lhs, scalar)");

EnvArguments 上的 SimpleOp

我们知道某些操作可能需要以下 −

标量作为输入，例如梯度标度
一组控制行为的关键字参数
用于加速计算的临时空间。

使用 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 { // 资源类型，指示指针类型是什么
        kRandom, // mshadow::Random<xpu> 对象
        kTempSpace // 可以是任意大小的动态临时空间
    };
    Type type; // 资源类型
};

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

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

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 是 x，b 是 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;
         }
      }
   };
}

Apache MXNet 教程

Apache MXNet 有用资源

Apache MXNet - 统一运算符 API

SimpleOp

定义形状

定义函数

定义梯度

将 SimpleOp 注册到 MXNet

EnvArguments 上的 SimpleOp

示例

构建张量操作

示例

颜色选择器

读后有收获微信请站长喝咖啡

错误报告

您的建议:

感谢您的帮助！