Apache MXNet - API d'opérateur unifié

Ce chapitre fournit des informations sur l'interface de programmation d'application (API) unifiée d'opérateur dans Apache MXNet.

SimpleOp

SimpleOp est une nouvelle API d'opérateur unifiée qui unifie différents processus d'appel. Une fois invoqué, il revient aux éléments fondamentaux des opérateurs. L'opérateur unifié est spécialement conçu pour les opérations unaires et binaires. C'est parce que la plupart des opérateurs mathématiques s'occupent d'un ou deux opérandes et que plusieurs opérandes rendent l'optimisation, liée à la dépendance, utile.

Nous allons comprendre son opérateur unifié SimpleOp fonctionnant à l'aide d'un exemple. Dans cet exemple, nous allons créer un opérateur fonctionnant comme unsmooth l1 loss, qui est un mélange de pertes l1 et l2. Nous pouvons définir et écrire la perte comme indiqué ci-dessous -

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

Ici, dans l'exemple ci-dessus,

. * représente la multiplication par élément
f, f’ est la fonction de perte l1 lisse dont nous supposons qu'elle est mshadow.

Il semble impossible d'implémenter cette perte particulière en tant qu'opérateur unaire ou binaire, mais MXNet fournit à ses utilisateurs une différenciation automatique en exécution symbolique qui simplifie directement la perte en f et f '. C'est pourquoi nous pouvons certainement implémenter cette perte particulière en tant qu'opérateur unaire.

Définition des formes

Comme nous le savons, MXNet mshadow librarynécessite une allocation de mémoire explicite, nous devons donc fournir toutes les formes de données avant tout calcul. Avant de définir les fonctions et le dégradé, nous devons fournir la cohérence de la forme d'entrée et la forme de sortie comme suit:

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 fonction mxnet :: Tshape est utilisée pour vérifier la forme des données d'entrée et la forme des données de sortie désignées. Dans ce cas, si vous ne définissez pas cette fonction, la forme de sortie par défaut serait la même que la forme d'entrée. Par exemple, dans le cas d'un opérateur binaire, la forme de lhs et rhs est cochée par défaut comme identique.

Passons maintenant à notre smooth l1 loss example. Pour cela, nous devons définir un XPU vers cpu ou gpu dans l'implémentation de l'en-tête smooth_l1_unary-inl.h. La raison est de réutiliser le même code dans smooth_l1_unary.cc et smooth_l1_unary.cu.

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

Comme dans notre smooth l1 loss example,la sortie a la même forme que la source, nous pouvons utiliser le comportement par défaut. Il peut s'écrire comme suit -

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

Définition des fonctions

Nous pouvons créer une fonction unaire ou binaire avec une entrée comme suit -

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

Voici le RunContext ctx struct qui contient les informations nécessaires lors de l'exécution pour l'exécution -

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

Voyons maintenant comment nous pouvons écrire les résultats du calcul dans 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
};

Maintenant, passons à notre smooth l1 loss example. Pour cela, nous utiliserons UnaryFunction pour définir la fonction de cet opérateur comme suit:

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

Définition des dégradés

Sauf Input, TBlob, et OpReqTypesont doublées, les fonctions Gradients des opérateurs binaires ont une structure similaire. Voyons ci-dessous, où nous avons créé une fonction de dégradé avec différents types d'entrée:

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

Tel que défini ci-dessus Input0, Input, OutputValue, et OutputGrad tous partagent la structure de GradientFunctionArgument. Il est défini comme suit -

struct GradFunctionArgument {
   TBlob data;
}

Passons maintenant à notre smooth l1 loss example. Pour que cela active la règle de la chaîne du gradient, nous devons multiplierout_grad du haut au résultat 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)));
   });
}

Enregistrer SimpleOp sur MXNet

Une fois que nous avons créé la forme, la fonction et le dégradé, nous devons les restaurer à la fois dans un opérateur NDArray et dans un opérateur symbolique. Pour cela, nous pouvons utiliser la macro d'enregistrement comme suit -

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

le SimpleOpInplaceOption peut être défini comme suit -

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

Passons maintenant à notre smooth l1 loss example. Pour cela, nous avons une fonction de gradient qui repose sur des données d'entrée afin que la fonction ne puisse pas être écrite en place.

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

Comme nous le savons, certaines opérations peuvent nécessiter les éléments suivants:

Un scalaire comme entrée telle qu'une échelle de gradient
Un ensemble d'arguments de mots clés contrôlant le comportement
Un espace temporaire pour accélérer les calculs.

L'avantage d'utiliser EnvArguments est qu'il fournit des arguments et des ressources supplémentaires pour rendre les calculs plus évolutifs et plus efficaces.

Exemple

Définissons d'abord la structure comme ci-dessous -

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

Ensuite, nous devons demander des ressources supplémentaires telles que mshadow::Random<xpu> et espace mémoire temporaire de EnvArguments.resource. Cela peut être fait comme suit -

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

Maintenant, l'enregistrement demandera la demande de ressource déclarée à partir de mxnet::ResourceManager. Après cela, il placera les ressources dans std::vector<Resource> resource in EnvAgruments.

Nous pouvons accéder aux ressources à l'aide du code suivant -

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

Si vous voyez dans notre exemple de perte l1 lisse, une entrée scalaire est nécessaire pour marquer le point de retournement d'une fonction de perte. C'est pourquoi dans le processus d'inscription, nous utilisonsset_enable_scalar(true), et env.scalar dans les déclarations de fonction et de gradient.

Opération Tensor de bâtiment

Ici, la question se pose de savoir pourquoi nous devons élaborer des opérations tensorielles? Les raisons sont les suivantes -

Le calcul utilise la bibliothèque mshadow et nous n'avons parfois pas de fonctions facilement disponibles.
Si une opération n'est pas effectuée de manière élémentaire telle que la perte et le gradient softmax.

Exemple

Ici, nous utilisons l'exemple de perte l1 lisse ci-dessus. Nous allons créer deux mappeurs, à savoir les cas scalaires de perte et de gradient l1 lisses:

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