Apache MXNet - API nhà điều hành hợp nhất

Chương này cung cấp thông tin về giao diện lập trình ứng dụng điều hành hợp nhất (API) trong Apache MXNet.

SimpleOp

SimpleOp là một API toán tử hợp nhất mới hợp nhất các quy trình gọi khác nhau. Sau khi được gọi, nó quay trở lại các phần tử cơ bản của các toán tử. Toán tử hợp nhất được thiết kế đặc biệt cho các phép toán đơn phân cũng như nhị phân. Đó là bởi vì hầu hết các toán tử toán học tham gia vào một hoặc hai toán hạng và nhiều toán hạng hơn làm cho việc tối ưu hóa, liên quan đến phụ thuộc, hữu ích.

Chúng ta sẽ hiểu toán tử hợp nhất SimpleOp của nó hoạt động với sự trợ giúp của một ví dụ. Trong ví dụ này, chúng tôi sẽ tạo một toán tử hoạt động như mộtsmooth l1 loss, là hỗn hợp mất đi l1 và l2. Chúng ta có thể xác định và viết khoản lỗ như dưới đây:

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

Đây, trong ví dụ trên,

  • . * là viết tắt của phép nhân theo nguyên tố

  • f, f’ là hàm mất mát l1 trơn mà chúng ta đang giả định là mshadow.

Có vẻ như không thể thực hiện sự mất mát cụ thể này dưới dạng toán tử một ngôi hoặc nhị phân nhưng MXNet cung cấp cho người dùng sự khác biệt tự động trong việc thực thi biểu tượng giúp đơn giản hóa sự mất mát đối với f và f 'một cách trực tiếp. Đó là lý do tại sao chúng tôi chắc chắn có thể thực hiện tổn thất cụ thể này với tư cách là toán tử một ngôi.

Xác định hình dạng

Như chúng ta biết MXNet's mshadow libraryyêu cầu phân bổ bộ nhớ rõ ràng do đó chúng tôi cần cung cấp tất cả các hình dạng dữ liệu trước khi xảy ra bất kỳ phép tính nào. Trước khi xác định các hàm và gradient, chúng ta cần cung cấp tính nhất quán của hình dạng đầu vào và hình dạng đầu ra như sau:

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

Hàm mxnet :: Tshape được sử dụng để kiểm tra hình dạng dữ liệu đầu vào và hình dạng dữ liệu đầu ra được chỉ định. Trong trường hợp, nếu bạn không xác định hàm này thì hình dạng đầu ra mặc định sẽ giống như hình dạng đầu vào. Ví dụ, trong trường hợp toán tử nhị phân, hình dạng của lhs và rhs theo mặc định được kiểm tra là giống nhau.

Bây giờ chúng ta hãy chuyển sang smooth l1 loss example. Đối với điều này, chúng ta cần xác định XPU thành cpu hoặc gpu trong triển khai tiêu đề smooth_l1_unary-inl.h. Lý do là để sử dụng lại cùng một mã trong smooth_l1_unary.ccsmooth_l1_unary.cu.

#include <mxnet/operator_util.h>
   #if defined(__CUDACC__)
      #define XPU gpu
   #else
      #define XPU cpu
#endif

Như trong của chúng tôi smooth l1 loss example,đầu ra có cùng hình dạng với nguồn, chúng ta có thể sử dụng hành vi mặc định. Nó có thể được viết như sau:

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

Xác định chức năng

Chúng ta có thể tạo một hàm đơn phân hoặc nhị phân với một đầu vào như sau:

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

Sau đây là RunContext ctx struct chứa thông tin cần thiết trong thời gian chạy để thực thi -

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

Bây giờ, hãy xem cách chúng ta có thể viết kết quả tính toán trong 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
};

Bây giờ, hãy chuyển sang smooth l1 loss example. Đối với điều này, chúng tôi sẽ sử dụng UnaryFunction để xác định chức năng của toán tử này như sau:

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

Định nghĩa Gradients

Ngoại trừ Input, TBlob,OpReqTypeđược nhân đôi, các hàm Gradients của toán tử nhị phân có cấu trúc tương tự. Hãy kiểm tra bên dưới, nơi chúng tôi đã tạo một hàm gradient với nhiều loại đầu vào:

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

Như đã định nghĩa ở trên Input0, Input, OutputValue,OutputGrad tất cả đều chia sẻ cấu trúc của GradientFunctionArgument. Nó được định nghĩa như sau:

struct GradFunctionArgument {
   TBlob data;
}

Bây giờ chúng ta hãy chuyển sang smooth l1 loss example. Đối với điều này để kích hoạt quy tắc chuỗi của gradient, chúng ta cần nhânout_grad từ đầu đến kết quả của 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)));
   });
}

Đăng ký SimpleOp với MXNet

Khi chúng ta đã tạo hình dạng, hàm và gradient, chúng ta cần khôi phục chúng thành cả toán tử NDArray cũng như thành toán tử tượng trưng. Đối với điều này, chúng ta có thể sử dụng macro đăng ký như sau:

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

Các SimpleOpInplaceOption có thể được định nghĩa như sau:

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

Bây giờ chúng ta hãy chuyển sang smooth l1 loss example. Đối với điều này, chúng ta có một hàm gradient dựa vào dữ liệu đầu vào để hàm không thể được viết tại chỗ.

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 trên EnvArguments

Như chúng ta biết một số thao tác có thể cần những điều sau:

  • Một đại lượng vô hướng như đầu vào chẳng hạn như thang độ dốc

  • Một tập hợp các đối số từ khóa kiểm soát hành vi

  • Một không gian tạm thời để tăng tốc độ tính toán.

Lợi ích của việc sử dụng EnvArguments là nó cung cấp các đối số và tài nguyên bổ sung để làm cho các phép tính có thể mở rộng và hiệu quả hơn.

Thí dụ

Đầu tiên, hãy xác định cấu trúc như bên dưới:

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

Tiếp theo, chúng tôi cần yêu cầu các tài nguyên bổ sung như mshadow::Random<xpu> và không gian bộ nhớ tạm thời từ EnvArguments.resource. Nó có thể được thực hiện như sau:

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

Bây giờ, đăng ký sẽ yêu cầu yêu cầu tài nguyên được khai báo từ mxnet::ResourceManager. Sau đó, nó sẽ đặt các tài nguyên vào std::vector<Resource> resource in EnvAgruments.

Chúng tôi có thể truy cập các tài nguyên với sự trợ giúp của mã sau:

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

Nếu bạn thấy trong ví dụ về mất mát l1 trơn tru của chúng tôi, thì cần có một đầu vào vô hướng để đánh dấu bước ngoặt của một hàm lỗ. Đó là lý do tại sao trong quá trình đăng ký, chúng tôi sử dụngset_enable_scalar(true)env.scalar trong khai báo hàm và gradient.

Xây dựng hoạt động Tensor

Ở đây câu hỏi đặt ra là tại sao chúng ta cần phải tạo ra các hoạt động tensor? Lý do như sau:

  • Tính toán sử dụng thư viện mshadow và đôi khi chúng tôi không có sẵn các chức năng.

  • Nếu một hoạt động không được thực hiện theo cách khôn ngoan như mất softmax và gradient.

Thí dụ

Ở đây, chúng tôi đang sử dụng ví dụ mất mát l1 mịn ở trên. Chúng tôi sẽ tạo hai ánh xạ cụ thể là các trường hợp vô hướng của mất mát l1 và gradient:

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