Apache MXNet - API operatore unificato

Questo capitolo fornisce informazioni sull'API (Application Programming Interface) per operatore unificato in Apache MXNet.

SimpleOp

SimpleOp è una nuova API operatore unificata che unifica diversi processi di chiamata. Una volta richiamato, ritorna agli elementi fondamentali degli operatori. L'operatore unificato è appositamente progettato per operazioni unarie e binarie. È perché la maggior parte degli operatori matematici si occupa di uno o due operandi e più operandi rendono utile l'ottimizzazione relativa alla dipendenza.

Comprenderemo il suo operatore unificato SimpleOp lavorando con l'aiuto di un esempio. In questo esempio, creeremo un operatore che funziona come filesmooth l1 loss, che è una combinazione di perdita l1 e l2. Possiamo definire e scrivere la perdita come indicato di seguito:

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

Qui, nell'esempio sopra,

  • . * sta per moltiplicazione basata sugli elementi

  • f, f’ è la funzione di perdita l1 regolare in cui supponiamo sia mshadow.

Sembra impossibile implementare questa particolare perdita come operatore unario o binario, ma MXNet fornisce ai suoi utenti una differenziazione automatica in esecuzione simbolica che semplifica direttamente la perdita in f ed f '. Ecco perché possiamo certamente implementare questa particolare perdita come operatore unario.

Definizione di forme

Come sappiamo di MXNet mshadow libraryrichiede un'allocazione esplicita della memoria, quindi dobbiamo fornire tutte le forme dei dati prima che avvenga qualsiasi calcolo. Prima di definire funzioni e gradiente, è necessario fornire la coerenza della forma di input e la forma di output come segue:

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 funzione mxnet :: Tshape viene utilizzata per controllare la forma dei dati di input e la forma dei dati di output designata. Nel caso in cui, se non si definisce questa funzione, la forma di output predefinita sarebbe la stessa della forma di input. Ad esempio, in caso di operatore binario, la forma di lhs e rhs è selezionata per impostazione predefinita come la stessa.

Ora passiamo al nostro smooth l1 loss example. Per questo, dobbiamo definire un XPU su cpu o gpu nell'implementazione dell'intestazione smooth_l1_unary-inl.h. Il motivo è riutilizzare lo stesso codice in 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

Come nel nostro smooth l1 loss example,l'uscita ha la stessa forma della sorgente, possiamo usare il comportamento predefinito. Può essere scritto come segue:

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

Definizione di funzioni

Possiamo creare una funzione unaria o binaria con un input come segue:

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

Di seguito è riportato il file RunContext ctx struct che contiene le informazioni necessarie durante il runtime per l'esecuzione -

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

Vediamo ora come scrivere i risultati del calcolo 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
};

Ora passiamo al nostro smooth l1 loss example. Per questo, useremo UnaryFunction per definire la funzione di questo operatore come segue:

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

Definizione dei gradienti

Tranne Input, TBlob, e OpReqTypesono raddoppiati, le funzioni Gradients degli operatori binari hanno una struttura simile. Diamo un'occhiata di seguito, dove abbiamo creato una funzione gradiente con vari tipi di input:

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

Come definito sopra Input0, Input, OutputValue, e OutputGrad tutti condividono la struttura di GradientFunctionArgument. È definito come segue:

struct GradFunctionArgument {
   TBlob data;
}

Ora passiamo al nostro smooth l1 loss example. Per abilitare la regola della catena del gradiente, dobbiamo moltiplicareout_grad dall'alto al risultato di 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)));
   });
}

Registra SimpleOp su MXNet

Dopo aver creato la forma, la funzione e il gradiente, è necessario ripristinarli sia in un operatore NDArray che in un operatore simbolico. Per questo, possiamo utilizzare la macro di registrazione come segue:

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

Il SimpleOpInplaceOption può essere definito come 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)
};

Ora passiamo al nostro smooth l1 loss example. Per questo, abbiamo una funzione gradiente che si basa sui dati di input in modo che la funzione non possa essere scritta sul posto.

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

Come sappiamo, alcune operazioni potrebbero richiedere quanto segue:

  • Uno scalare come input come una scala gradiente

  • Un insieme di argomenti di parole chiave che controllano il comportamento

  • Uno spazio temporaneo per velocizzare i calcoli.

Il vantaggio dell'utilizzo di EnvArguments è che fornisce argomenti e risorse aggiuntivi per rendere i calcoli più scalabili ed efficienti.

Esempio

Per prima cosa definiamo la struttura come di seguito -

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

Successivamente, dobbiamo richiedere risorse aggiuntive come mshadow::Random<xpu> e spazio di memoria temporaneo da EnvArguments.resource. Può essere fatto come segue:

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

Ora, la registrazione richiederà la richiesta di risorsa dichiarata da mxnet::ResourceManager. Dopodiché, inserirà le risorse std::vector<Resource> resource in EnvAgruments.

Possiamo accedere alle risorse con l'aiuto del seguente codice:

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

Se nel nostro esempio di perdita l1 regolare, è necessario un input scalare per contrassegnare il punto di svolta di una funzione di perdita. Ecco perché nel processo di registrazione utilizziamoset_enable_scalar(true), e env.scalar nelle dichiarazioni di funzione e gradiente.

Building Tensor Operation

Qui sorge la domanda: perché dobbiamo creare operazioni tensoriali? I motivi sono i seguenti:

  • Il calcolo utilizza la libreria mshadow ea volte non abbiamo funzioni prontamente disponibili.

  • Se un'operazione non viene eseguita in un modo saggio per elementi come softmax loss e gradiente.

Esempio

Qui, stiamo usando l'esempio di perdita l1 liscia sopra. Creeremo due mappatori, vale a dire i casi scalari di perdita l1 liscia 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;
         }
      }
   };
}