Apache MXNet - унифицированный API оператора
В этой главе содержится информация об унифицированном интерфейсе прикладного программирования (API) оператора в Apache MXNet.
SimpleOp
SimpleOp - это новый унифицированный API-интерфейс оператора, который объединяет различные процессы вызова. После вызова он возвращается к основным элементам операторов. Унифицированный оператор специально разработан как для унарных, так и для бинарных операций. Это потому, что большинство математических операторов относятся к одному или двум операндам, а большее количество операндов делает оптимизацию, связанную с зависимостями, полезной.
Разберемся его унифицированный оператор SimpleOp, работающий на примере. В этом примере мы создадим оператор, работающий какsmooth l1 loss, который представляет собой смесь потерь 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 libraryтребует явного выделения памяти, поэтому нам нужно предоставить все формы данных до того, как произойдет какое-либо вычисление. Перед определением функций и градиента нам необходимо обеспечить согласованность формы ввода и формы вывода следующим образом:
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 loss example. Для этого нам нужно определить XPU для cpu или gpu в реализации заголовка. smooth_l1_unary-inl.h. Причина в том, чтобы повторно использовать тот же код в 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 example,выход имеет ту же форму, что и исходный, мы можем использовать поведение по умолчанию. Это можно записать следующим образом -
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 который содержит информацию, необходимую во время выполнения для выполнения -
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
};А теперь перейдем к нашему smooth l1 loss example. Для этого мы будем использовать 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;
}А теперь перейдем к нашему smooth l1 loss example. Для этого, чтобы включить цепное правило градиента, нам нужно умножить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)
};А теперь перейдем к нашему smooth l1 loss example. Для этого у нас есть функция градиента, которая полагается на входные данные, поэтому функцию нельзя записать на месте.
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
};Затем нам нужно запросить дополнительные ресурсы, например mshadow::Random<xpu> и временное пространство памяти из EnvArguments.resource. Это можно сделать следующим образом -
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. После этого он разместит ресурсы в std::vector<Resource> resource in 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 loss и gradient.
пример
Здесь мы используем приведенный выше пример сглаживания потерь 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;
}
}
};
}