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.