MxNet学习: 深度学习的依赖引擎

版权声明:本文为博主原创文章,未经博主允许不得转载。 https://blog.csdn.net/cdknight_happy/article/details/84621824

我们总是希望深度学习库运行得更快,能够扩展到更大的数据集。一种自然的方法是,看看我们是否可以从向这个问题投入更多硬件中获益,就像同时使用多个gpu。

然后,库设计人员会问:如何跨设备并行计算?更重要的是,我们如何在引入多线程时同步计算?运行时依赖引擎是这些问题的通用解决方案。

在本文中,我们研究了使用运行时依赖调度来加速深度学习的方法。我们的目的是解释运行时依赖调度如何既能加速又能简化多设备深度学习。我们还探讨了一种通用依赖引擎的潜在设计,这种引擎既可以是独立于库的,也可以是独立于操作的。

本文的大部分讨论都是从MXNet依赖引擎中获得灵感的。我们讨论的依赖跟踪算法最初是由Yutian Li和Mingjie Wang开发的。

1 依赖调度

虽然大多数用户希望利用并行计算,但我们大多数人更熟悉串行程序。因此,一个自然的问题是:我们如何编写串行程序并构建一个库以异步方式自动并行我们的程序?

例如,在下面的代码中,我们可以按任意顺序运行B = A + 1和C = A + 2,或者并行运行:

    A = 2
    B = A + 1
    C = A + 2
    D = B * C

但是,手工编写序列代码是相当困难的,因为最后一个操作D = B * C需要等待前面的两个操作都完成之后才能开始。下面的依赖关系图/数据流图说明了这一点。
在这里插入图片描述
依赖引擎是一个库,它接受一系列操作,并根据依赖模式(可能是并行的)对它们进行调度。因此在本例中,依赖项库可以并行运行B = a + 1和C = a + 2,并在这些操作完成后运行D = B * C。

2 依赖调度中的问题

依赖引擎减轻了编写并发程序的负担。然而,随着操作的并行化,新的依赖跟踪问题出现了。在本节中,我们将讨论这些问题。

2.1 数据流依赖

数据流依赖关系描述如何在其他计算中使用一个计算的结果。每个依赖引擎都必须解决数据流依赖问题。
在这里插入图片描述
因为我们在前一节中讨论了这个问题,所以这里包含了相同的图。拥有数据流跟踪引擎的库包括Minerva和Purine2。

2.2 内存复用

什么时候应该回收分配给数组的内存?在串行处理中,这很容易确定。我们只是在变量超出作用域后回收内存。然而,如下图所示,这在并行处理中有点困难。
在这里插入图片描述
在这个例子中,因为这两个计算都需要使用来自A的值,所以在完成之前我们不能回收内存。引擎必须根据依赖关系调度内存回收操作,并确保在B = A + 1和C = A + 2完成之后执行这些操作。

2.3 随机数生成

随机数生成器是机器学习中常用的一种生成器,它给依赖引擎带来了有趣的挑战。考虑下面的例子:
在这里插入图片描述
在这个例子中,我们按顺序生成随机数。虽然这两个随机数生成似乎可以并行化,但通常不是这样。伪随机数生成器(PRNG)不是线程安全的,因为在生成新数字时,它可能导致某些内部状态发生突变。即使PRNG是线程安全的,序列化数字生成也是可取的,这样我们就可以得到可重复的随机数。

2.4 案例研究:一个多gpu神经网络的依赖引擎

在上一节中,我们讨论了在设计依赖引擎时可能遇到的问题。在考虑如何设计一个通用引擎来解决这些问题之前,让我们先考虑一个依赖引擎如何在神经网络的多gpu训练中提供帮助。下面的伪代码Python程序演示了在双层神经网络上的一个批次的训练过程。

    # Example of one iteration Two GPU neural Net
    data = next_batch()
    data[gpu0].copyfrom(data[0:50])
    data[gpu1].copyfrom(data[50:100])
    # forward, backprop on GPU 0
    fc1[gpu0] = FullcForward(data[gpu0], fc1_weight[gpu0])
    fc2[gpu0] = FullcForward(fc1[gpu0], fc2_weight[gpu0])
    fc2_ograd[gpu0] = LossGrad(fc2[gpu0], label[0:50])
    fc1_ograd[gpu0], fc2_wgrad[gpu0] =
      FullcBackward(fc2_ograd[gpu0] , fc2_weight[gpu0])
      _, fc1_wgrad[gpu0] = FullcBackward(fc1_ograd[gpu0] , fc1_weight[gpu0])
    # forward, backprop on GPU 1
    fc1[gpu1] = FullcForward(data[gpu1], fc1_weight[gpu1])
    fc2[gpu1] = FullcForward(fc1[gpu1], fc2_weight[gpu1])
    fc2_ograd[gpu1] = LossGrad(fc2[gpu1], label[50:100])
    fc1_ograd[gpu1], fc2_wgrad[gpu1] =
         FullcBackward(fc2_ograd[gpu1] , fc2_weight[gpu1])
         _, fc1_wgrad[gpu1] = FullcBackward(fc1_ograd[gpu1] , fc1_weight[gpu1])
    # aggregate gradient and update
    fc1_wgrad[cpu]  = fc1_wgrad[gpu0] + fc1_wgrad[gpu1]
    fc2_wgrad[cpu]  = fc2_wgrad[gpu0] + fc2_wgrad[gpu1]
    fc1_weight[cpu] -= lr *  fc1_wgrad[cpu]
    fc2_weight[cpu] -= lr *  fc2_wgrad[cpu]
    fc1_weight[cpu].copyto(fc1_weight[gpu0] , fc1_weight[gpu1])
    fc2_weight[cpu].copyto(fc2_weight[gpu0] , fc2_weight[gpu1])

在这个程序中,数据0到50复制到GPU 0,数据50到100复制到GPU 1。计算出的梯度在CPU中聚合,然后CPU执行一个简单的SGD更新,并将更新后的权重复制回每个GPU。这是一种以串行方式编写并行程序的常见方法。下图显示了如何并行化它:
在这里插入图片描述
说明:

  • 一旦我们得到一个层的梯度,可以立即将其复制到CPU;
  • 权重一经更新即可复制回GPU;
  • 在前向计算中,有一个对上一轮迭代的依赖项:fc1_weight[cpu].copyto(fc1_weight[gpu0] , fc1_weight[gpu1]);
  • 从上一次向后传递到第k层到下一次向前调用第k层之间有一个计算延迟,在此延迟期间,我们可以将k层的权值同步与其他计算并行进行同步。

这种优化方法被CXXNet等多gpu深度学习库所使用。重点是使权重同步(通信)与计算重叠。但是,这并不容易,因为复制操作需要在层的后向传递完成后立即触发,然后再触发reduce、update等。

依赖引擎可以调度这些操作并执行多线程和依赖跟踪。

3 设计一个通用的依赖引擎

我们希望您确信依赖引擎对于将深度学习程序扩展到多个设备是有用的。现在让我们讨论如何为依赖引擎设计和实现通用接口。这种解决方案并不是依赖引擎唯一可能的设计。这是一个我们认为在大多数情况下有用的例子。

我们的目标是创建一个通用的轻量级依赖引擎。理想情况下,我们希望引擎能够方便地插入到现有的深度学习代码中,并且能够在稍加修改的情况下扩展到多台机器。要做到这一点,我们只需要关注依赖跟踪,而不是假设用户可以或不能做什么。

这里是引擎目标的总结:

  • 引擎不应该知道它执行什么操作,以便用户可以执行他们定义的任何操作;
  • 不应该限制引擎可以调度的对象类型;
    我们应该能够调度对GPU和CPU内存的依赖 ;
    我们应该能够跟踪随机数生成器的依赖关系,等等;
  • 引擎不应该分配资源。它应该只跟踪依赖项。用户可以分配自己的内存、PRNG等。

下面的Python代码片段提供了一个引擎接口,可以帮助我们实现目标。注意,真正的实现将更接近于metal,通常是在c++中。

    class DepEngine(object):
        def new_variable():
            """Return a new variable tag
            Returns
            -------
            vtag : Variable Tag
                The token of the engine to represent dependencies.
            """
            pass

        def push(exec_func, read_vars, mutate_vars):
            """Push the operation to the engine.

            Parameters
            ----------
            exec_func : callable
                The real operation to be performed.

            read_vars : list of Variable Tags
                The list of variables this operation will read from.

            mutate_vars : list of Variable Tags
                The list of variables this operation will mutate.
            """
            pass

因为我们不能假设我们正在调度什么对象,所以我们要求用户分配一个与每个对象相关联的虚拟标记来表示我们需要调度什么。因此,在一开始,用户可以分配变量标记,并将其附加到我们希望调度的每个对象上。
在这里插入图片描述
然后用户调用push来告诉引擎要执行的函数。用户还需要指定操作的依赖关系,使用read_vars和write_vars:

  • read_vars是只读变量的标记,不需要更改其内部状态;
  • mutate_vars是内部状态可改变的对象的变量标记。
    在这里插入图片描述上面的图显示了如何将操作B = A + 1推送到依赖引擎。B.data和A.data是分配好的空间。请注意,引擎只知道变量标记。任何执行函数都可以被处理。这个接口对于我们想要调度的操作和资源是通用的。

为了好玩,让我们通过以下代码片段来了解引擎内部是如何处理标记的:

    B = A + 1
    C = A + 2
    A = C * 2
    D = A + 3

第一行读取变量A和可变变量B。第二行读取变量A和可变变量C,以此类推。

引擎为每个变量维护一个队列,如下面的动画所示。绿色块表示读取操作,而红色块表示修改操作。

在这里插入图片描述在构建此队列时,引擎会看到A队列开头的前两个绿色块实际上可以并行运行,因为它们都是读取操作,不会相互冲突。下图说明了这一点。

在这里插入图片描述
所有这些调度的一个很酷的地方是,它不只局限于数值计算。因为计划的所有事情都只是一个标记,引擎可以计划所有事情!

下图给出了我们在前面几节中提到的程序的完整推送序列。
在这里插入图片描述

3.1 将现有代码移植到依赖引擎

因为通用接口不控制内存分配和要执行的操作,所以依赖引擎可以将大多数现有代码调度为两个步骤:

  • 分配与内存blob、PRNG等资源相关的变量标记;
  • 调用push,将执行函数作为要执行的原始代码,并将相应资源的变量标记正确地放入read_vars和mutate_vars中。

3.2 实现通用依赖引擎

我们已经描述了通用引擎接口,以及如何使用它来调度各种操作。在本节中,我们将对如何实现这样的引擎进行高级讨论。

总体思路如下:

  • 使用队列跟踪每个变量标记上所有挂起的依赖项;
  • 在每个操作上使用计数器来跟踪还需要完成多少依赖项;
  • 操作完成后,更新队列和依赖项计数器的状态,以调度新的操作。

下图演示了调度算法,可能会让您更好地了解引擎中正在发生的事情。
在这里插入图片描述
下面,我们将展示另一个涉及随机数生成器的示例。
在这里插入图片描述
如您所见,该算法的目的是更新操作的挂起队列,并在操作完成时进行正确的状态转换。应该更加小心地确保状态转换以线程安全的方式完成。

3.3 独立的依赖跟踪和运行策略

如果您正在仔细阅读,您可能已经注意到前面的部分只显示了决定何时可以执行操作的算法。我们没有展示如何实际运行一个操作。实际上,可以有许多不同的策略。例如,我们可以使用一个全局线程池来运行所有操作,或者使用一个特定的线程在每个设备上运行操作。

这种运行策略通常独立于依赖跟踪,可以分为独立模块或基本依赖跟踪模块的虚拟接口。开发一个对所有操作和调度都公平的优雅的运行时策略本身就是一个有趣的系统问题。

4 讨论

我们在本文中讨论的设计并不是依赖跟踪问题的唯一解决方案。这只是我们如何解决这个问题的一个例子。当然,其中一些设计选择是有争议的。我们将在本节中讨论其中一些。

4.1 动态 VS. 静态

本主题中讨论的依赖项引擎接口在某种程度上是动态的,因为用户可以逐个推进操作,而不是声明整个依赖项图(静态)。就数据结构而言,动态调度可能比静态声明需要更多的开销。但是,它也支持更多的灵活性,例如支持命令式程序的自动并行性,或者命令式程序和符号程序的混合。您还可以向接口添加一些预先声明的操作,以支持数据结构重用。

4.2 可变的 VS. 不可变

本文提供的通用引擎接口支持显式的可变调度。在典型的数据流引擎中,数据通常是不可变的。使用不可变数据有很多好处。例如,不可变数据通常更适合并行化,并且有助于在分布式设置中更好地容错(通过重新计算)。

然而,不可变性带来了几个挑战:

  • 在处理随机数和删除时,很难安排资源竞争问题;
  • 引擎通常需要管理资源(内存、随机数)以避免冲突。很难处理用户分配的空间,等等;
  • 预先分配的静态内存不可用,这也是因为通常的模式是写入预先分配的层空间,如果数据是不可变的,则不支持这种模式。

允许可变调度可以缓解这些问题。

4.3 MxNet

MxNet是按照本文的思想实现了通用的依赖引擎。

5 参考

文章翻译自:https://mxnet.incubator.apache.org/architecture/note_engine.html

展开阅读全文

没有更多推荐了,返回首页