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函数中，我们可以传递线程的大部分，但是执行引擎直到我们调用回调函数才认为该函数已完成。

语境

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

该函数是否应在 CPU 或 GPU 上运行。
如果我们在Context中指定GPU，那么就使用哪个GPU。
Context 和 RunContext 之间存在巨大差异。Context 具有设备类型和设备 ID，而 RunContext 具有只能在运行时确定的信息。

变量句柄

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

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

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

它是轻量级的，因此创建、删除或复制变量只需要很少的运营成本。
我们需要指定不可变变量，即将在const_vars中使用的变量。
我们需要指定可变变量，即将在mutate_vars中修改的变量。
执行引擎解决函数之间依赖关系的规则是，当任意两个函数之一修改至少一个公共变量时，它们的执行将按照其推送顺序进行序列化。
为了创建新变量，我们可以使用NewVar() API。
为了删除变量，我们可以使用PushDelete API。

让我们通过一个简单的例子来理解它的工作原理 -

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

推动并等待

推和等待是执行引擎的两个比较有用的API。

以下是Push API的两个重要特性：

所有 Push API 都是异步的，这意味着无论推送的函数是否完成，API 调用都会立即返回。
Push API 不是线程安全的，这意味着一次只能有一个线程进行引擎 API 调用。

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

如果用户想要等待特定函数完成，他/她应该在闭包中包含回调函数。包含后，在函数末尾调用该函数。
另一方面，如果用户想要等待涉及某个变量的所有函数完成，他/她应该使用WaitForVar(var) API。
如果有人想等待所有推送的函数完成，请使用WaitForAll () API。
用于指定函数的依赖关系，就像一个令牌。

运营商

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

操作界面

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

各种任务

操作员界面允许用户执行以下任务 -

用户可以指定就地更新并可以减少内存分配成本
为了使其更简洁，用户可以隐藏 Python 的一些内部参数。
用户可以定义张量和输出张量之间的关系。
为了执行计算，用户可以从系统获取额外的临时空间。

操作符属性

我们知道，在卷积神经网络（CNN）中，一个卷积有多种实现方式。为了从它们中获得最佳性能，我们可能需要在这几个卷积之间进行切换。

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

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

第一个目的是告诉系统每个输入和输出张量的大小，以便在前向和后向调用之前分配空间。
第二个目的是在运行之前执行大小检查以确保没有错误。

语法如下 -

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

这里，需要注意的两个要点 -

FulllyConnectedForward中的out_data不被FullyConnectedBackward使用，并且
PoolingBackward 需要 PoolingForward 的所有参数。

这就是为什么对于FullyConnectedForward，一旦消耗了out_data张量就可以安全地释放，因为后向函数不需要它。在这个系统的帮助下，我们可以尽早收集一些张量作为垃圾。

In place Option - 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

创建算子

首先在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步

参数化运算符

如果要实现卷积运算符，则必须知道内核大小、步长大小、填充大小等。为什么，因为这些参数应该在调用任何前向或后向接口之前传递给操作员。

为此，我们需要定义一个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

将运算符属性类和参数类注册到 Apache MXNet

最后，我们需要将 Operator Property Class 和 Parameter Class 注册到 MXNet。可以借助以下宏来完成 -

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

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