---

Apache MXNet - 系统组件

---

这里详细介绍了 Apache MXNet 中的系统组件。首先，我们将研究 MXNet 中的执行引擎。

执行引擎

Apache MXNet 的执行引擎非常通用。我们可以将它用于深度学习以及任何特定领域的问题：按照它们的依赖关系执行一组函数。它的设计方式是，具有依赖关系的函数被序列化，而没有依赖关系的函数可以并行执行。

核心接口

下面给出的 API 是 Apache MXNet 执行引擎的核心接口 −

virtual void PushSync(Fn exec_fun, Context exec_ctx,
std::vector<VarHandle> const& const_vars,
std::vector<VarHandle> const& mutate_vars) = 0;

上述 API 具有以下 −

• exec_fun − MXNet 的核心接口 API 允许我们将名为 exec_fun 的函数连同其上下文信息和依赖项一起推送到执行引擎。

• exec_ctx − 应在其中执行上述函数 exec_fun 的上下文信息。

• const_vars − 这些是函数读取的变量。

• mutate_vars −这些是要修改的变量。

执行引擎为其用户提供保证，即修改公共变量的任何两个函数的执行都按其推送顺序进行序列化。

函数

以下是 Apache MXNet − 执行引擎的函数类型

using Fn = std::function<void(RunContext)>;

在上述函数中，RunContext 包含运行时信息。运行时信息应由执行引擎确定。RunContext 的语法如下 −

struct RunContext {
   // stream pointer which could be safely cast to
   // cudaStream_t* type
   void *stream;
};

下面给出了有关执行引擎函数的一些要点−

• 所有函数均由 MXNet 执行引擎的内部线程执行。

• 将函数阻塞推给执行引擎并不好，因为这样函数将占用执行线程，还会降低总吞吐量。

为此，MXNet 提供了另一个异步函数，如下所示−

using Callback = std::function<void()>;
using AsyncFn = std::function<void(RunContext, Callback)>;

• 在这个 AsyncFn 函数中，我们可以传递线程的繁重部分，但执行引擎不会认为该函数已完成，直到我们调用 callback 函数。

Context

在 Context 中，我们可以指定要在其中执行的函数的上下文。这通常包括以下 −

• 该函数应该在 CPU 还是 GPU 上运行。

• 如果我们在 Context 中指定 GPU，那么应该使用哪个 GPU。

• Context 和 RunContext 之间存在巨大差异。 Context 有设备类型和设备 id，而 RunContext 有只能在运行时决定的信息。

VarHandle

VarHandle 用于指定函数的依赖关系，就像一个 token（特别是由执行引擎提供的），我们可以用它来表示函数可以修改或使用的外部资源。

但问题来了，为什么我们需要使用 VarHandle？这是因为，Apache MXNet 引擎被设计为与其他 MXNet 模块解耦。

以下是有关 VarHandle − 的一些要点

• 它是轻量级的，因此创建、删除或复制变量几乎不会产生任何操作成本。

• 我们需要指定不可变变量，即将在 const_vars 中使用的变量。

• 我们需要指定可变变量，即将在 mutate_vars 中修改的变量。

• 执行引擎用于解决函数间依赖关系的规则是，当任何两个函数中的一个修改至少一个公共变量时，它们的执行将按照它们的推送顺序进行序列化。

• 要创建新变量，我们可以使用NewVar() API。

• 要删除变量，我们可以使用 PushDelete API。

让我们通过一个简单的示例了解其工作原理 −

假设我们有两个函数，即 F1 和 F2，它们都改变变量，即 V2。在这种情况下，如果 F2 在 F1 之后被推送，则保证 F2 在 F1 之后执行。另一方面，如果 F1 和 F2 都使用 V2，那么它们的实际执行顺序可能是随机的。

推送和等待

推送 和 等待 是执行引擎另外两个有用的 API。

以下是 推送 API 的两个重要功能：

• 所有 Push API 都是异步的，这意味着无论推送的函数是否完成，API 调用都会立即返回。

• Push API 不是线程安全的，这意味着一次只能有一个线程进行引擎 API 调用。

现在，如果我们谈论 Wait API，以下几点代表它 −

• 如果用户想要等待特定函数完成，他/她应该在闭包中包含一个回调函数。包含后，在函数末尾调用该函数。

• 另一方面，如果用户想要等待涉及某个变量的所有函数完成，则应使用 WaitForVar(var) API。

• 如果有人想等待所有推送的函数完成，则使用 WaitForAll () API。

• 用于指定函数的依赖关系，就像一个令牌。

运算符

Apache MXNet 中的运算符是一个包含实际计算逻辑以及辅助信息并帮助系统执行优化的类。

运算符接口

Forward 是核心运算符接口，其语法如下：

virtual void Forward(const OpContext &ctx,
const std::vector<TBlob> &in_data,
const std::vector<OpReqType> &req,
const std::vector<TBlob> &out_data,
const std::vector<TBlob> &aux_states) = 0;

在Forward()中定义的OpContext的结构如下：

struct OpContext {
   int is_train;
   RunContext run_ctx;
   std::vector<Resource> requested;
}

OpContext描述了算子的状态（是在训练阶段还是测试阶段）、算子应该在哪个设备上运行以及请求的资源。执行引擎另外两个有用的API。

从上面的Forward核心接口，我们可以理解请求的资源如下 −

• in_data和out_data代表输入和输出张量。

• req表示计算结果如何写入out_data。

OpReqType可以定义为 −

enum OpReqType {
   kNullOp,
   kWriteTo,
   kWriteInplace,
   kAddTo
};

与 Forward 运算符一样，我们可以选择性地实现 Backward 接口，如下所示 −

virtual void Backward(const OpContext &ctx,
const std::vector<TBlob> &out_grad,
const std::vector<TBlob> &in_data,
const std::vector<TBlob> &out_data,
const std::vector<OpReqType> &req,
const std::vector<TBlob> &in_grad,
const std::vector<TBlob> &aux_states);

各种任务

Operator 接口允许用户执行以下任务 −

• 用户可以指定就地更新并减少内存分配成本

• 为了使其更简洁，用户可以隐藏 Python 中的一些内部参数。

• 用户可以定义张量和输出张量之间的关系。

• 要执行计算，用户可以从系统获取额外的临时空间。

Operator 属性

众所周知，在卷积神经网络 (CNN) 中，一个卷积有几种实现。为了从中获得最佳性能，我们可能需要在这几种卷积之间切换。

这就是 Apache MXNet 将运算符语义接口与实现接口分开的原因。这种分离以 OperatorProperty 类的形式完成，该类由以下内容组成−

InferShape − InferShape 接口有两个用途，如下所示：

• 第一个用途是告诉系统每个输入和输出张量的大小，以便在 Forward 和 Backward 调用之前分配空间。

• 第二个用途是执行大小检查，以确保在运行之前没有错误。

语法如下 −

virtual bool InferShape(mxnet::ShapeVector *in_shape,
mxnet::ShapeVector *out_shape,
mxnet::ShapeVector *aux_shape) const = 0;

请求资源 − 如果您的系统可以管理 cudnnConvolutionForward 等操作的计算工作区，会怎么样？您的系统可以执行优化，例如重用空间等等。在这里，MXNet 借助以下两个接口轻松实现这一点 −

virtual std::vector<ResourceRequest> ForwardResource(
   const mxnet::ShapeVector &in_shape) const;
virtual std::vector<ResourceRequest> BackwardResource(
   const mxnet::ShapeVector &in_shape) const;

但是，如果 ForwardResource 和 BackwardResource 返回非空数组怎么办？在这种情况下，系统通过 Operator 的 Forward 和 Backward 接口中的 ctx 参数提供相应的资源。

向后依赖 − Apache MXNet 有以下两种不同的运算符签名来处理向后依赖 −

void FullyConnectedForward(TBlob weight, TBlob in_data, TBlob out_data);
void FullyConnectedBackward(TBlob weight, TBlob in_data, TBlob out_grad, TBlob in_grad);
void PoolingForward(TBlob in_data, TBlob out_data);
void PoolingBackward(TBlob in_data, TBlob out_data, TBlob out_grad, TBlob in_grad);

这里要注意两个要点 −

• FullyConnectedForward 中的 out_data 未被 FullyConnectedBackward 使用，并且

• PoolingBackward 需要 PoolingForward 的所有参数。

这就是为什么对于 FullyConnectedForward，out_data 张量一旦被使用就可以被安全释放，因为后向函数将不需要它。借助此系统，可以尽早将一些张量收集为垃圾。

Inplace 选项 − Apache MXNet 为用户提供了另一个接口，以节省内存分配成本。该接口适用于输入和输出张量具有相同形状的元素级操作。

以下是指定就地更新 − 的语法

创建运算符的示例

借助 OperatorProperty，我们可以创建一个运算符。为此，请按照以下步骤操作−

virtual std::vector<std::pair<int, void*>> ElewiseOpProperty::ForwardInplaceOption(
   const std::vector<int> &in_data,
   const std::vector<void*> &out_data) 
const {
   return { {in_data[0], out_data[0]} };
}
virtual std::vector<std::pair<int, void*>> ElewiseOpProperty::BackwardInplaceOption(
   const std::vector<int> &out_grad,
   const std::vector<int> &in_data,
   const std::vector<int> &out_data,
   const std::vector<void*> &in_grad) 
const {
   return { {out_grad[0], in_grad[0]} }
}

步骤 1

创建 Operator

首先在 OperatorProperty 中实现以下接口：

virtual Operator* CreateOperator(Context ctx) const = 0;

示例如下 −

class ConvolutionOp {
   public:
      void Forward( ... ) { ... }
      void Backward( ... ) { ... }
};
class ConvolutionOpProperty : public OperatorProperty {
   public:
      Operator* CreateOperator(Context ctx) const {
         return new ConvolutionOp;
      }
};

步骤 2

参数化运算符

如果您要实现卷积运算符，则必须知道内核大小、步幅大小、填充大小等。这是因为在调用任何 Forward 或 backward 接口之前，应将这些参数传递给运算符。

为此，我们需要定义一个 ConvolutionParam 结构，如下所示−

#include <dmlc/parameter.h>
struct ConvolutionParam : public dmlc::Parameter<ConvolutionParam> {
   mxnet::TShape kernel, stride, pad;
   uint32_t num_filter, num_group, workspace;
   bool no_bias;
};

现在，我们需要将其放入 ConvolutionOpProperty 中，并将其传递给运算符，如下所示 −

class ConvolutionOp {
   public:
      ConvolutionOp(ConvolutionParam p): param_(p) {}
      void Forward( ... ) { ... }
      void Backward( ... ) { ... }
   private:
      ConvolutionParam param_;
};
class ConvolutionOpProperty : public OperatorProperty {
   public:
      void Init(const vector<pair<string, string>& kwargs) {
         // initialize param_ using kwargs
      }
      Operator* CreateOperator(Context ctx) const {
         return new ConvolutionOp(param_);
      }
   private:
      ConvolutionParam param_;
};

步骤 3

将 Operator Property Class 和 Parameter Class 注册到 Apache MXNet

最后，我们需要将 Operator Property Class 和 Parameter Class 注册到 MXNet。这可以通过以下宏来完成 −

DMLC_REGISTER_PARAMETER(ConvolutionParam);
MXNET_REGISTER_OP_PROPERTY(Convolution, ConvolutionOpProperty);

在上面的宏中，第一个参数是名称字符串，第二个参数是属性类名。

---