Apache MXNet - API Unified Operator

Este capítulo fornece informações sobre a interface de programação de aplicativo (API) do operador unificada no Apache MXNet.

SimpleOp

SimpleOp é uma nova API de operador unificada que unifica diferentes processos de invocação. Depois de invocado, ele retorna aos elementos fundamentais dos operadores. O operador unificado é especialmente projetado para operações unárias e binárias. É porque a maioria dos operadores matemáticos atende a um ou dois operandos e mais operandos tornam a otimização, relacionada à dependência, útil.

Estaremos entendendo seu operador unificado SimpleOp trabalhando com a ajuda de um exemplo. Neste exemplo, estaremos criando um operador funcionando como umsmooth l1 loss, que é uma mistura de perda de l1 e l2. Podemos definir e escrever a perda conforme abaixo -

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

Aqui, no exemplo acima,

  • . * significa multiplicação elementar

  • f, f’ é a função de perda suave l1 que estamos assumindo está em mshadow.

Parece impossível implementar essa perda em particular como um operador unário ou binário, mas o MXNet fornece a seus usuários diferenciação automática na execução simbólica que simplifica a perda para f e f 'diretamente. É por isso que podemos certamente implementar essa perda específica como um operador unário.

Definindo formas

Como sabemos do MXNet mshadow libraryrequer alocação de memória explícita, portanto, precisamos fornecer todas as formas de dados antes que qualquer cálculo ocorra. Antes de definir funções e gradiente, precisamos fornecer consistência da forma de entrada e forma de saída da seguinte maneira:

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

A função mxnet :: Tshape é usada para verificar o formato dos dados de entrada e o formato dos dados de saída designados. No caso, se você não definir esta função, o formato de saída padrão será o mesmo que o formato de entrada. Por exemplo, no caso do operador binário, a forma de lhs e rhs é marcada como igual por padrão.

Agora vamos passar para o nosso smooth l1 loss example. Para isso, precisamos definir uma XPU para cpu ou gpu na implementação do cabeçalho smooth_l1_unary-inl.h. O motivo é reutilizar o mesmo código em smooth_l1_unary.cc e smooth_l1_unary.cu.

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

Como em nosso smooth l1 loss example,a saída tem a mesma forma da fonte, podemos usar o comportamento padrão. Pode ser escrito da seguinte forma -

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

Definindo Funções

Podemos criar uma função unária ou binária com uma entrada da seguinte maneira -

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

A seguir está o RunContext ctx struct que contém as informações necessárias durante o tempo de execução para execução -

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

Agora, vamos ver como podemos escrever os resultados do cálculo em 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
};

Agora, vamos passar para o nosso smooth l1 loss example. Para isso, usaremos UnaryFunction para definir a função deste operador da seguinte maneira:

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

Definindo Gradientes

Exceto Input, TBlob, e OpReqTypesão duplicados, as funções de gradientes de operadores binários têm uma estrutura semelhante. Vamos verificar a seguir, onde criamos uma função gradiente com vários tipos de entrada:

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

Conforme definido acima Input0, Input, OutputValue, e OutputGrad todos compartilham a estrutura de GradientFunctionArgument. É definido como segue -

struct GradFunctionArgument {
   TBlob data;
}

Agora vamos passar para o nosso smooth l1 loss example. Para que isso habilite a regra da cadeia de gradiente, precisamos multiplicarout_grad do topo ao resultado de 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)));
   });
}

Registre SimpleOp para MXNet

Depois de criar a forma, a função e o gradiente, precisamos restaurá-los tanto em um operador NDArray quanto em um operador simbólico. Para isso, podemos usar a macro de registro da seguinte forma -

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

o SimpleOpInplaceOption pode ser definido como segue -

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

Agora vamos passar para o nosso smooth l1 loss example. Para isso, temos uma função gradiente que depende dos dados de entrada para que a função não possa ser escrita no local.

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 em EnvArguments

Como sabemos, algumas operações podem precisar do seguinte -

  • Um escalar como entrada, como uma escala de gradiente

  • Um conjunto de argumentos de palavras-chave controlando o comportamento

  • Um espaço temporário para agilizar os cálculos.

A vantagem de usar EnvArguments é que ele fornece argumentos e recursos adicionais para tornar os cálculos mais escalonáveis ​​e eficientes.

Exemplo

Primeiro, vamos definir a estrutura conforme abaixo -

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

Em seguida, precisamos solicitar recursos adicionais, como mshadow::Random<xpu> e espaço de memória temporário de EnvArguments.resource. Isso pode ser feito da seguinte forma -

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

Agora, o registro irá solicitar a solicitação de recurso declarada de mxnet::ResourceManager. Depois disso, ele colocará os recursos em std::vector<Resource> resource in EnvAgruments.

Podemos acessar os recursos com a ajuda do seguinte código -

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

Se você ver em nosso exemplo de perda suave l1, uma entrada escalar é necessária para marcar o ponto de viragem de uma função de perda. É por isso que, no processo de registro, usamosset_enable_scalar(true), e env.scalar em declarações de função e gradiente.

Operação do tensor de construção

Aqui surge a pergunta: por que precisamos criar operações de tensor? As razões são as seguintes -

  • A computação utiliza a biblioteca mshadow e às vezes não temos funções disponíveis.

  • Se uma operação não for realizada de maneira elementar, como softmax loss e gradiente.

Exemplo

Aqui, estamos usando o exemplo de perda suave l1 acima. Estaremos criando dois mapeadores, a saber, os casos escalares de perda suave l1 e gradiente:

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