Dive into Deep

MXNet 系统架构（翻译）

翻译自 https://mxnet.incubator.apache.org/architecture/overview.html

mxnet architecture overview

这张图展示了 MXNet 的主要模块和组件，以及它们之间的交互。这些模块是：

MXNet 系统模块

Execution Engine (执行引擎)

你不仅可以使用 MXNet 引擎进行深度学习，还可以使用它来解决任何专业领域的问题。MXNet 的设计目标是解决通用问题：根据依赖关系来执行一系列的函数。有依赖关系的任意两个函数应该按顺序执行。为提升性能，没有依赖关系的函数可以被并行执行。关于这方面的更详细的讨论，请见 notes on the dependency engine。

Interface (接口)

下面这个 API 是执行引擎的核心接口：

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）和它的执行上下文（context）以及依赖关系一起推送给执行引擎。exec_ctx 是 exec_fun 的执行上下文，const_vars 是函数要读取的变量，mutate_vars 是函数要修改的变量。执行引擎提供以下保证：

任意两个函数，如果它们要修改同一个变量，则它们的执行顺序与它们被推送到引擎的顺序是一致的。

Function (函数)

引擎的函数的类型是：

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

RunContext 包含了运行时的信息，这些信息由引擎来确定。

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

或者，你也可以使用 mxnet::engine::DAGEngine::Fn, 它有相同的类型定义。

所有的函数都在引擎的内部线程中执行。在这种模型中，把会阻塞的函数（通常是处理磁盘或网络之类的 I/O 任务的操作）发送给引擎通常不是个好主意，因为它会占用执行线程，并且降低吞吐量。在这种情况下，我们提供了另一个异步的函数：

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

在 AsyncFn 中，你可以把阻塞的部分传回到你自己的线程，然后退出函数。引擎在调用 Callback 之后才会认为 AsyncFn 函数执行完成。

Context (上下文)

你可以指定函数执行的上下文。这通常包括函数该在 CPU 还是 GPU 上运行，如果指定使用 GPU，还可以指定哪个 GPU。Context 和 RunContext 不同，Context 包含设备类型 (GPU/CPU) 和设备 id，而 RunContext 包含只有运行时才能确定的信息，比如函数应该在哪个流上执行。

VarHandle

VarHandle 是用来指定函数之间的依赖关系的。MXNet 引擎被设计为与其他模块之间是松耦合的，因此 VarHandle 就像一个引擎提供的标记，你可以用它来代表函数能使用或修改的外部资源。VarHandle 很轻量，所以创建、删除或复制都只有很小的开销。在把函数推送给引擎时，你需要在 const_vars 向量中指定函数将要使用（只读）的变量，在 mutate_vars 向量中指定函数将要修改的变量。引擎使用以下规则来解析函数之间的依赖关系：

如果两个函数修改至少一个共同的变量，那么它们的执行顺序和它们被推送的顺序是一致的。

举个例子，如果 Fn1 和 Fn2 都要修改 V2，并且 Fn2 是在 Fn1 之后被推送的，则执行引擎保证 Fn2 在 Fn1 之后执行。如果 Fn1 和 Fn2 都使用（只读）V2，则它们的执行顺序是随机的。

这样的设计允许引擎最小化内存分配的方式来调度运算。例如，DNN 的权重更新函数可以使用 += 操作来进行原地操作，而不是每次都生成新的数组。

你可以使用 NewVar() 来创建变量，用 PushDelete() 来删除变量。

Push and Wait (推送和等待)

所有的 Push API 都是异步的，函数调用会立即返回，不管 Fn 是否执行完成。这允许引擎在用户线程推送函数的时候并行开始执行计算。Push API 不是线程安全的，在同一个时刻，只有一个线程可以调用 Push API。

如果你想要等待某个特定的 Fn 执行完成，你可以传入一个 Callback，并且在你的 Fn 结束的时候调用它。

如果你想等待与某个变量相关的所有 Fn 都结束，可以使用 WaitForVar(var) API。

如果你想等待所有已推送的 Fn 都结束，可以使用 WaitForAll() API。

Save Object Creation Cost (减少创建对象的开销)

在某些情况下，你需要在较长一段时间内把多个函数推送到引擎。如果这些函数的计算量不大，那么复制匿名函数和创建变量的开销就会变得相对比较高。这种情况下，我们提供了 API 来提前创建 OprHandle：

virtual OprHandle NewOperator(AsyncFn fn,
                              std::vector<VarHandle> const& const_vars,
                              std::vector<VarHandle> const& mutate_vars) = 0;

你可以连续推送 OprHandle 而不用重复创建它们：

virtual void Push(OprHandle op, Context exec_ctx) = 0;

要删除它，调用 DeleteOperator(OprHandle op) API。要保证在调用这个 API 之前，运算符已经完成计算。

API Reference

略。

Operators (运算符) in MXNet

在 MXNet 中，运算符是一个类，包括了实际的计算逻辑和一些帮助系统进行优化的辅助信息，比如原地更新和自动求导之类的。要理解这篇文档余下的部分，我们建议你熟悉一下 mshadow 库，因为所有的运算符都是在系统运行时提供的类似张量（tensor-like）的数据结构 mshadow::TBlob 上进行计算。

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;

OpContext 结构体：

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

它描述了此运算符是在训练阶段还是测试阶段，它应该运行在哪个设备上（run_ctx），以及已经请求的资源（这个会在之后的章节讨论）。

除了 Foward 操作，你可以视需要选择实现 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);

这个接口遵循与 Forward 相同的设计原则，不同之处是接口中 out_grad, in_data 和 out_data 是给定的，需要计算 in_grad 作为结果。这里的命名规则与 Torch 类似，可以总结为下图：

[input/output semantics figure]

某些运算符可能可以省略某个参数：out_grad, in_data, 和 out_data。你可以用 OperatorProperty 类的 DeclareBackwardDependency 接口来指定它们的依赖关系。

Operator Property (运算符属性)

卷积可以有多种实现方式，你可能想在各种方式之间切换以实现最佳性能。因此，我们把运算符的语义接口从他的实现接口（Operator 类）中剥离出来，放到 OperatorProperty 类中。OperatorProperty 接口包含：

从 Operator Property 创建 Operator

OperatorProperty 包含了 Operator 的所有的语义属性。它也负责为实际计算创建 Operator。

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

Operator 的参数化

当你实现一个卷积运算符时，你需要知道核的大小 (kernel size)，步长的大小 (stride size)，填充的大小 (padding size)，等等。这些参数应当在 Forward 和 Backward 接口被调用之前传给 Operator。你可以定义一个 ConvolutionParam 结构，如下：

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

把它放在 ConvolutionOpProperty 中，并且在创建 Operator 时传进去：

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

使用以下宏来把 Operator 的 Property 类和 Parameter 类注册到 MXNet：

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

第一个参数是名字，第二个参数是 Property 的类名。

接口总结

到目前为止，我们基本上涵盖了定义一个 Operator 的全部接口。让我们回顾一下：

统一 NDArray 运算符和符号运算符

NDArray 运算符和符号运算符类似，区别是在没有完整的依赖关系图时，有时你不能原地更新。然而，NDArray 运算符和符号运算符的底层逻辑是一样的。SimpleOp 是一个新的统一的运算符 API，统一了不同的方式。因为多数数学运算符有一个或两个操作数，而更多个操作数使得和依赖关系相关的优化有用，统一的运算符是被专门设计用于一元和二元运算。

让我们考虑运算符的基本元素。理想情况下，你只需要用函数和导数来描述一个运算符。让我们将讨论限制在一元和二元运算符。我们该如何分类所有操作符，来使原地更新的可能性最大？注意你可以按照操作数的数量来对函数进行分类。而导数更复杂一些。要创建一个依赖关系图，你需要知道输出值，输入数据是否在之后的梯度中被用到。统一 API 中的梯度函数是用操作数的类型来区分的。

在我们继续了解 SimpleOp 接口之前，我们建议你浏览 mshadow library guide，因为计算是在 mshadow:TBlob 中进行的。

在以下例子中，我们会创建一个实现 smooth l1 loss 函数的运算符，这是 l1 loss 和 l2 loss 的混合体。这个函数可以被写成如下形式：

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

.* 代表元素乘，f 和 f’ 是 smooth l1 loss 函数，我们先假设它们在 mshadow 中可以找到。乍看起来，不可能把它实现成一元或二元操作。但是我们有符号的自动微分，这简化了 f 和 f’ 的损失函数。这个损失函数并不比 sin 或者 abs 函数复杂，我们可以将它实现为一元操作符。

SimpleOp: 统一的 Operator API

Define Shapes (定义形状)

mshadow 库需要显式地分配内存，因此在计算开始前，就需要定义好所有数据的形状。在定义函数和导数之前，先让我们检查输入的形状并且提供输出的形状。

···cpp
typedef TShape (UnaryShapeFunction)(const TShape& src,const EnvArguments& env);
typedef TShape (BinaryShapeFunction)(const TShape& const lhs, TShape& rhs, const EnvArguments& env);
···

你可以使用 mshadow::TShape 来检查输入形状并且指定输出形状。如果你不定义这个函数，则默认的输出形状和输入形状相同。如果是二元运算符，默认情况下，系统会检查 lhs 和 rhs 的形状是否相同。

你还可以用这些函数来检查是否有额外的参数和资源。可以参考 EnvArguments 的用法。

在我们开始 smooth l1 loss 的例子之前，我们在头文件 smooth_l1_unary-inl.h 中定义了 XPU，值为 cpu 或者 gpu，以便于我们可以在 smooth_l1_unary.cc 和 smoth_l1_unary.cu 中使用相同的代码。

#include
#if defined(__CUDACC__)
#define XPU gpu
#else
#define XPU cpu
#endif

在 smooth l1 loss 的例子中，因为输入输出有相同的形状，我们就直接用默认行为：

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

定义函数

创建有一个输出 (mshadow::TBlob) 的一元或二元函数。

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

在 smooth l1 loss 例子中，我们使用 UnaryFunction 来定义操作符的函数。

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

从 RunContext 获取到 mshadow::Stream 之后，我们从 mshadow::TBlob 拿到 mshadow::Tensor。mshadow::F 是创建一个 mshadow 表达式的快捷方式。宏 MSHADOW_TYPE_SWITCH(type, DType, …) 处理不同类型的细节，宏 ASSIGN_DISPATCH(out, req, exp) 检查 OpReqType 并且执行相应动作。 sigma2 是这个损失函数中的一个特殊的参数，我们后面会讲到。

定义梯度（可选）

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

二元运算符的梯度函数拥有相似的结构，不同之处在于 Input, TBlob 和 OpReqType 的数量翻倍。

GradFunctionArgument

Input0, Input, OutputValue 和 OutputGrad 都和 GradFunctionArgument 有相同的结构，定义如下：

struct GradFunctionArgument {
    TBlob data;
}

在 smooth l1 loss 例子中，注意用到输入来计算梯度的是 f’(x)，所以 UnaryGradFunctionT2 是适用的。我们还需要用 out_grade 乘以结果的 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)));
    });
}

把 SimpleOp 注册到 MXNet

在创建 shape, function 和 gradient 之后，要把它们放到 NDArray 运算符和符号运算符中。可以使用 operator_util.h 中定义的宏来简化这个过程。

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

SimpleOpInplaceOption 的定义：

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

在我们的例子中，我们的梯度依赖于函数的输入，因此函数不能原地更新数据。输出的梯度在梯度计算后就不被用到了，因此梯度可以原地更新。

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

不要忘了之前的讨论，在没有用 set_shape_function 来设置 shape 的时候，默认会强制输出的 shape 和输入的 shape 一致。我们后面会讨论 set_enable_scalar。

NDArray 运算符总结

SimpleOp 的其他信息

在 EnvArguments 上使用 SimpleOp

一些操作需要一个标量作为输入，比如一个梯度标量，一组用来控制算法行为的关键字参数，或者一个用来加速计算的临时空间。EnvArguments 提供额外的参数和资源来使计算更易扩展和更高效。

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

要启用这些额外的功能，需要更多的注册参数。为避免参数的混淆，scalar 和 kwargs 不能同时使用。要启用 scalar，在注册时使用 set_enable_scalar(bool enable_scalar)。之后，在前向函数和梯度中，可以用函数测参数 EnvArguments env 中的 env.scalar 来访问 scalar。

要启用 kwargs，在注册时使用 set_enable_kwargs(bool enable_kwargs)。在前向函数和反向梯度时，额外的参数包含在 env.kwargs 中，被定义为 std::vectorstd::string。使用 DMLC 参数接口来简化关键字参数的解析。更多细节请参考 guide on parameter structure。

mshadow::Random 或者临时内存空间之类的额外资源可以通过 EnvArguments.resoure 来请求和访问。注册方式为 set_resource_request(ResourceRequest req) 或者 set_resource_request(const std::vector)，其中 mxnet::ResourceRequest 被定义为：

struct ResourceRequest {
    enum Type {  // Resource type, indicating what the pointer type is
        kRandom,  // mshadow::Random object
        kTempSpace  // A dynamic temp space that can be arbitrary size
    };
    Type type;  // type of resources
};

相关例子请参见 src/operator/loss_binary_op-inl.h。

在我们的 smooth l1 loss 例子中，需要有一个标量输入来标记损失函数的折点。因此，在注册时，我们使用 set_enable_scalar(true)，并且在函数和梯度中使用 env.scalar。

实现一个张量运算 (Tensor Operation)

因为使用 mshadow 库来进行计算，有时候没有我们用到的函数，我们可以在运算符中实现张量运算。如果你把函数定义为元素 (element-wise) 运算，那么你可以把它实现成一个 mxnet::op::mshadow_op。src/operator/mshadow_op.h 中有许多 mshadow_op 的例子。 mshadow_op 是表达式映射器。它们处理函数的标量形式。细节请见 mshadow expression API guide。

如果一个运算不能用元素 (element-wise) 方式实现，比如 softmax 损失函数和梯度，那么你就需要实现一个新的张量运算。你需要直接创建 mshadow 函数和 mshadow::cuda 函数。更多例子请见 src/ooperator/roi_pooling.cc。

在我们的 smooth l1 loss 例子中，我们创建了两个映射器，分别是标量情形下的 smooth l1 loss 和 gradient。

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

梯度 (gradient) 与之相似，可以在 src/operator/smooth_l1_unary-inl.h 中找到。

两个以上操作数

新的统一 API 被设计为完成运算符的基本功能。对于有两个以上输入、或一个以上输出、或是需要更多特性的运算符，请见原始的 Operator API。

标签

akka architecture deeplearning haskell java linux math monad mxnet nginx php probability python pytorch statistics

归档

2020 Jan 2017 Nov 2017 Sep 2017 Jun 2017 Feb 2016 Nov 2016 Oct 2015 Mar 2014 Oct 2013 Jul 2013 Jun 2013 May 2009 Jan