Apache MXNet: API de operador unificado

Este capítulo proporciona información sobre la interfaz de programación de aplicaciones (API) del operador unificada en Apache MXNet.

SimpleOp

SimpleOp es una nueva API de operador unificada que unifica diferentes procesos de invocación. Una vez invocado, vuelve a los elementos fundamentales de los operadores. El operador unificado está especialmente diseñado para operaciones unarias y binarias. Es porque la mayoría de los operadores matemáticos atienden a uno o dos operandos y más operandos hacen que la optimización, relacionada con la dependencia, sea útil.

Comprenderemos su operador unificado SimpleOp trabajando con la ayuda de un ejemplo. En este ejemplo, crearemos un operador que funcione comosmooth l1 loss, que es una mezcla de pérdida de l1 y l2. Podemos definir y escribir la pérdida como se indica a continuación:

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

Aquí, en el ejemplo anterior,

. * significa multiplicación por elementos
f, f’ es la función de pérdida suave l1 que asumimos está en mshadow.

Parece imposible implementar esta pérdida en particular como un operador unario o binario, pero MXNet proporciona a sus usuarios una diferenciación automática en la ejecución simbólica que simplifica la pérdida af y f 'directamente. Es por eso que ciertamente podemos implementar esta pérdida en particular como un operador unario.

Definición de formas

Como sabemos de MXNet mshadow libraryrequiere una asignación de memoria explícita, por lo que debemos proporcionar todas las formas de datos antes de que se produzca cualquier cálculo. Antes de definir funciones y gradientes, debemos proporcionar consistencia de forma de entrada y forma de salida de la siguiente manera:

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

La función mxnet :: Tshape se usa para verificar la forma de los datos de entrada y la forma de los datos de salida designados. En caso de que no defina esta función, la forma de salida predeterminada sería la misma que la forma de entrada. Por ejemplo, en el caso de un operador binario, la forma de lhs y rhs está marcada por defecto como la misma.

Ahora pasemos a nuestro smooth l1 loss example. Para esto, necesitamos definir una XPU a cpu o gpu en la implementación del encabezado smooth_l1_unary-inl.h. La razón es reutilizar el mismo código en smooth_l1_unary.cc y smooth_l1_unary.cu.

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

Como en nuestro smooth l1 loss example,la salida tiene la misma forma que la fuente, podemos usar el comportamiento predeterminado. Se puede escribir de la siguiente manera:

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

Definición de funciones

Podemos crear una función unaria o binaria con una entrada de la siguiente manera:

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

Lo siguiente es el RunContext ctx struct que contiene la información necesaria durante el tiempo de ejecución para la ejecución -

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

Ahora, veamos cómo podemos escribir los resultados del cálculo en 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
};

Ahora, pasemos a nuestro smooth l1 loss example. Para ello, usaremos UnaryFunction para definir la función de este operador de la siguiente manera:

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

Definición de degradados

Excepto Input, TBlob, y OpReqTypese duplican, las funciones de degradado de los operadores binarios tienen una estructura similar. Veamos a continuación, donde creamos una función de gradiente con varios 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);

Como se define arriba Input0, Input, OutputValue, y OutputGrad todos comparten la estructura de GradientFunctionArgument. Se define de la siguiente manera:

struct GradFunctionArgument {
   TBlob data;
}

Ahora pasemos a nuestro smooth l1 loss example. Para que esto habilite la regla de la cadena del gradiente, necesitamos multiplicarout_grad desde la cima hasta el 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)));
   });
}

Registrar SimpleOp en MXNet

Una vez que creamos la forma, la función y el gradiente, necesitamos restaurarlos tanto en un operador NDArray como en un operador simbólico. Para esto, podemos usar la macro de registro de la siguiente manera:

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

los SimpleOpInplaceOption se puede definir de la siguiente manera:

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

Ahora pasemos a nuestro smooth l1 loss example. Para esto, tenemos una función de gradiente que se basa en los datos de entrada para que la función no se pueda escribir en su lugar.

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

Como sabemos, algunas operaciones pueden necesitar lo siguiente:

Un escalar como entrada, como una escala de gradiente
Un conjunto de argumentos de palabras clave que controlan el comportamiento.
Un espacio temporal para acelerar los cálculos.

El beneficio de usar EnvArguments es que proporciona argumentos y recursos adicionales para hacer los cálculos más escalables y eficientes.

Ejemplo

Primero definamos la estructura de la siguiente manera:

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

A continuación, necesitamos solicitar recursos adicionales como mshadow::Random<xpu> y espacio de memoria temporal de EnvArguments.resource. Se puede hacer de la siguiente manera:

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

Ahora, el registro solicitará la solicitud de recurso declarada de mxnet::ResourceManager. Después de eso, colocará los recursos en std::vector<Resource> resource in EnvAgruments.

Podemos acceder a los recursos con la ayuda del siguiente 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);

Si ve en nuestro ejemplo de pérdida l1 suave, se necesita una entrada escalar para marcar el punto de inflexión de una función de pérdida. Por eso en el proceso de registro utilizamosset_enable_scalar(true)y env.scalar en declaraciones de función y gradiente.

Operación del tensor del edificio

Aquí surge la pregunta de que ¿por qué tenemos que crear operaciones de tensores? Las razones son las siguientes:

La computación utiliza la biblioteca mshadow y, a veces, no tenemos funciones disponibles.
Si una operación no se realiza en forma de elementos, como pérdida y gradiente de softmax.

Ejemplo

Aquí, estamos usando el ejemplo de pérdida suave l1 anterior. Crearemos dos mapeadores, a saber, los casos escalares de pérdida y gradiente suaves de 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;
         }
      }
   };
}

↰ Previous page Next page ↱