Apache Mxnet 简明教程

Apache MXNet - System Components

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

Execution Engine

Apache MXNet 的执行引擎非常通用。它可以用深度学习或任何特定领域问题:执行一些函数,同时遵循这些函数的依赖关系。它的设计方式使依赖函数序列化,而没有依赖关系的函数可以并行执行。

Core Interface

下面给出的 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 具有以下内容:

  1. exec_fun −MXNet 的核心接口 API 允许我们会同其上下文信息和依赖关系将名为 exec_fun 的函数推送到执行引擎。

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

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

  4. mutate_vars −这些是将修改的变量。

执行引擎向其用户保证以按顺序推入的方式对修改通用变量的两个函数的执行进行序列化。

Function

以下为 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;
};

以下列出有关执行引擎函数的一些重要提示:

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

  2. 最好不要将阻塞函数推送到执行引擎,因为那样函数将占据执行线程,并且还会降低总吞吐率。

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

using Callback = std::function<void()>;
using AsyncFn = std::function<void(RunContext, Callback)>;
  1. 在这个 AsyncFn 函数中,我们可以传递线程的繁重部分,但是直到调用 callback 函数,执行引擎才认为该函数已完成。

Context

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

  1. 函数应在 CPU 还是 GPU 上运行。

  2. 如果我们在上下文中指定 GPU,则指定使用哪个 GPU。

  3. Context 和 RunContext 之间有很大区别。Context 具有设备类型和设备 ID,而 RunContext 具有只能在运行时确定的信息。

VarHandle

用于指定函数依赖关系的 VarHandle 就像一个标记(特别是由执行引擎提供的),它可用于表示函数可以修改或使用的外部资源。

但出现了问题,为什么我们需要使用 VarHandle?这是因为,Apache MXNet 引擎被设计为与其他 MXNet 模块分离。

以下是有关 VarHandle 的一些重要要点:

  1. 它很轻量,因此创建、删除或复制变量几乎不需要操作成本。

  2. 我们需要指定不可变变量,即将在 const_vars 中使用的变量。

  3. 我们需要指定可变变量,即将在 mutate_vars 中修改的变量。

  4. 执行引擎用来解析函数之间依赖关系的规则是,当其中一个函数修改至少一个公共变量时,对这两个函数的执行按其推送顺序进行序列化。

  5. 对于创建新变量,我们可以使用 NewVar() API。

  6. 对于删除变量,我们可以使用 PushDelete API。

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

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

Push and Wait

Pushwait 是执行引擎中另外两个有用的 API。

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

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

  2. 推送 API 不是线程安全的,这意味着一次只能有一个线程进行引擎 API 调用。

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

  1. 如果用户想等待一个特定的函数完成,他/她应该在闭包中包含一个回调函数。包含之后,在函数的末尾调用函数。

  2. 另一方面,如果用户想等待涉及某个变量的所有函数完成,他/她应该使用 WaitForVar(var) API。

  3. 如果某人想等待所有推入的函数完成,那么使用 WaitForAll () API。

  4. 用于指定函数的依赖项,就像一个令牌。

Operators

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

Operator Interface

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 的核心接口,我们可以理解请求的资源如下 −

  1. in_dataout_data 代表输入和输出张量。

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

Various tasks

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

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

  2. 为了使其更简洁,用户可以从 Python 中隐藏一些内部参数。

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

  4. 为了执行计算,用户可以从系统获取额外的临时空间。

Operator Property

我们知道在卷积神经网络 (CNN) 中,一个卷积有多种实现。为了从中获得最佳性能,我们可能希望在这些卷积之中进行切换。

这就是 Apache MXNet 将算子语义接口从实现接口中分离出来的原因。此分离以以下形式完成: OperatorProperty 类,它包含以下内容:−

InferShape - InferShape 接口有两个目的,如下所示:

  1. 第一个目的是告诉系统每个输入和输出张量的尺寸,以便可以在 ForwardBackward 调用之前分配空间。

  2. 第二个目的是执行大小检查,以确保在运行之前不存在错误。

语法如下所示:−

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

Request Resource - 如果您的系统可以管理像 cudnnConvolutionForward 这样的操作的计算工作空间会怎么样?您的系统可以执行优化,例如重用空间和更多内容。在这里,MXNet 在以下两个接口的帮助下轻松实现了这一点−

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

但是,如果 ForwardResourceBackwardResource 返回非空数组会怎么样?在这种情况下,系统通过 ForwardBackward 接口的 ctx 参数提供相应的资源 Operator

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

在这里,需要注意的两个重要点:

  1. FullyConnectedForward 中的 out_data 不被 FullyConnectedBackward 使用,

  2. PoolingBackward 需要 PoolingForward 的所有参数。

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

In place Option - Apache MXNet 为用户提供了另一个接口来节省内存分配的成本。此接口适用于输入和输出张量具有相同形状的逐元素运算。

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

Example for Creating an Operator

借助 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]} }
}

Step 1

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

Step 2

Parameterize Operator

如果你要实施一个卷积运算符,必须知道核大小、步幅大小、填充大小等。因为在调用任何 Forwardbackward 接口之前,应将这些参数传递给运算符。

为此,我们需要定义一个 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_;
};

Step 3

Register the Operator Property Class and the Parameter Class to Apache MXNet

最后,我们需要将运算符属性类和参数类注册到 MXNet。可以使用以下宏来完成此操作 −

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

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