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 中，我们可以指定要在其中执行函数的上下文。这通常包括以下内容：

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

  • 如果我们在 Context 中指定 GPU，则使用哪个 GPU。

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

VarHandle

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

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

以下是一些关于 VarHandle 的重要说明：

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

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

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

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

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

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

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

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

Push 和 Wait

Push 和 wait 是执行引擎的另外两个有用的 API。

以下是 Push 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 中的一些内部参数。

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

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

算子属性

我们知道在卷积神经网络 (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 张量，就可以安全地释放它，因为反向函数不需要它。借助这个系统，可以尽早收集一些张量作为垃圾。

就地选项 − 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

参数化算子

如果您要实现卷积算子，则必须知道内核大小、步幅大小、填充大小等。为什么？因为这些参数应该在调用任何 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

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

最后，我们需要将算子属性类和参数类注册到 MXNet。这可以使用以下宏来完成：

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

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

打印页面
上一页
下一页
广告