[关闭]
@atry 2017-11-16T10:06:20.000000Z 字数 9875 阅读 1339

神经网络与函数式编程(七)如何实现Monadic Deep Learning?

神经网络与函数式编程

在本系列先前的文章神经网络与函数式编程(四)Monadic Deep Learning中,我们学习了DeepLearning.scala如何实现多态函数,把各种不同的参数统一转为Do[Tape]来处理。在本篇文章中,我们将讨论DeepLearning.scala的反向传播机制,看看DeepLearning.scala如何实现Monadic风格的可差分计算图。


方案一:基于Tape的反向传播

首先我将简化一下DeepLearning.scala的实现,把涉及Monad的部分去掉,展示一下如何DeepLearning.scala内部的反向传播流程。

以一元线性回归模型为例,如果采用LAD惩罚函数,可以用DeepLearning.scala表达成这样:

  1. val weight = DoubleWeight(math.random)
  2. def f(x: Double): DoubleLayer = x * weight
  3. def lossFunction(x: Double, expectedY: Double): DoubleLayer = {
  4. abs(expectedY - f(x))
  5. }

给定带标记的数据样本xexpectedY,可以用train函数训练:

  1. val x: Double = ???
  2. val expectedY: Double = ???
  3. lossFunction(x, expectedY).train

其中,lossFunction(x, expectedY)是一个计算图DoubleLayer。在本系列前一篇文章,我们已经知道DoubleLayer可以forward转为Do[Tape[Double, Double]]Do是个Monad类型,我们暂时忽略,把DoubleLayer.forward简单的看成Tape[Double, Double]。简化后的Tape定义如下:

  1. case class Tape[Data, Delta](data: Data, backward: Delta => Unit)

Tape包含前向阶段的结果data,以及用来进行反向传播的backward函数。

比如train是反向传播的起点,可以实现成这样:

  1. trait DeepLearning[Differentiable] {
  2. type Data
  3. type Delta
  4. def forward(differentiable: Differentiable): Tape[Data, Delta]
  5. def train(lossFunction: DoubleLayer): Data = {
  6. val tape: Tape[Data, Delta] = lossFunction.forward
  7. tape.backward(1.0)
  8. tape.data
  9. }
  10. }

train用到的lossFunction展开后是abs(expectedY - x * weight)

// TODO: 此处缺一张图

其中可差分的abs实现如下

  1. def abs(operand0: DoubleLayer): DoubleLayer = {
  2. def forward: Tape[Double, Double] = {
  3. val tape = operand0.forward
  4. val data = math.abs(tape.data)
  5. val backward = { delta =>
  6. tape.backward(math.sign(tape.data) * delta)
  7. }
  8. Tape(data, backward)
  9. }
  10. DoubleLayer(forward)
  11. }

可以看到在前向阶段,absdataoperand0的data的绝对值。而在反向传播时,delta函数把导数经过计算后,传给operand0backward

abs用到的operand0则是可差分的减法,根据求导公式可以写成这样:

  1. def -(operand0: DoubleLayer, operand1: DoubleLayer): DoubleLayer = {
  2. def forward: Tape[Double, Double] = {
  3. val tape0 = operand0.forward
  4. val tape1 = operand1.forward
  5. val data = tape0.data - tape1.data
  6. val backward = { delta =>
  7. tape0.backward(delta)
  8. tape1.backward(-delta)
  9. }
  10. Tape(data, backward)
  11. }
  12. DoubleLayer(forward)
  13. }

注意,减法与abs不同,是个二元运算符,所以backward时需要分别反向传播到两个操作数的backward

然后,减法的operand1f(x),是个可差分的乘法:

  1. def *(operand0: DoubleLayer, operand1: DoubleLayer): DoubleLayer = {
  2. def forward: Tape[Double, Double] = {
  3. val tape0 = operand0.forward
  4. val tape1 = operand1.forward
  5. val data = tape0.data * tape1.data
  6. val backward = { delta =>
  7. tape0.backward(delta * tape1.data)
  8. tape1.backward(delta * tape0.data)
  9. }
  10. Tape(data, backward)
  11. }
  12. DoubleLayer(forward)
  13. }

最后,weight是个DoubleWeightDoubleWeightbackward会修改自己内部保存的值:

  1. trait DoubleWeight {
  2. var data: Double
  3. def forward: Tape[Double, Double] = {
  4. val backward = { delta =>
  5. data -= delta
  6. }
  7. Tape(data, backward)
  8. }
  9. }

到此为止,可以看到,当用户调用train时。
train内部首先会递归调用forward,在最深处由DoubleWeight取得data,然后一层一层执行前向阶段,返回一个Tape,其中包含前些阶段的执行结果。
然后,train内部会递归调用Tape上的backward,把导数反向传播,直到传到DoubleWeight上,令DoubleWeight得到训练。

方案二:支持粗粒度并行计算的异步计算图

到目前为止,我们实现的计算图是同步执行的。这意味着,计算图上的节点都不可能并行执行,即使这些节点之间并没有相互依赖关系。为了实现神经网络与函数式编程(四)Monadic Deep Learning中提到的粗粒度并行计算,我们需要让forward返回一个异步任务,而不应该阻塞。

  1. trait DeepLearning[Differentiable] {
  2. type Data
  3. type Delta
  4. def forward(differentiable: Differentiable): UnitContinuation[Tape[Data, Delta]]
  5. def train(lossFunction: DoubleLayer): UnitContinuation[Data] = monadic[UnitContinuation] {
  6. val tape = lossFunction.forward.each
  7. tape.backward(UnitContinuation.now(1.0)).each
  8. tape.data
  9. }
  10. }

以上代码中,我们引入了UnitContinuation来表示异步任务。train也做了修改,改用monadic风格的控制流。

相应的Tape.backward也需要改为异步结构。

  1. case class Tape[Data, Delta](data: Data, backward: UnitContinuation[Delta] => UnitContinuation[Unit])

我们用了一个小技巧,让backward接受UnitContinuation[Delta]而不是Delta参数。这样一来,由于UnitContinuation[Delta]是个异步操作而没有严格求值,那么,当backward触发时,backward的实现者如果是不包含任何权重的计算图,backward就可以跳过不执行,最终以自动实现了类似PyTorch中requires_grad的优化效果,而无需用户任何手动设置。

最后,每一个原子操作也改为monadic风格,比如乘法操作:

  1. def *(operand0: DoubleLayer, operand1: DoubleLayer): DoubleLayer = {
  2. def forward: UnitContinuation[Tape[Double, Double]] = {
  3. val Parallel(result) = (Parallel(operand0.forward) |@| Parallel(operand1.forward)) { (tape0, tape1) =>
  4. val data = tape0.data * tape1.data
  5. val backward = { deltaContinuation =>
  6. deltaContinuation.flatMap { delta =>
  7. val backward0 = tape0.backward(delta * tape1.data)
  8. val backward1 = tape1.backward(delta * tape0.data)
  9. val Parallel(backwardBoth) = Parallel(backward0) |@| Parallel(backward1) { (_, _) => }
  10. backwardBoth
  11. }
  12. }
  13. Tape(data, backward)
  14. }
  15. result
  16. }
  17. DoubleLayer(forward)
  18. }

注意,我们实现了乘法这样的二元操作符时,无论是前向阶段还是反向阶段,都利用Parallel让两个操作数上的计算得到并行执行。

到此为止,方案二通过引入UnitContinuation,让所有计算图都得到了异步计算的能力,并且可以并行执行多个计算图的分支。

问题:菱形计算图导致指数时间复杂度

autograd等早期的自动求导算法不同,DeepLearning.scala在二元操作处,能对两个操作数都求导。你可以在减法和乘法的上看到,backward内部会分别触发两个操作数的backward。这种做法可以支持复杂的网络结构.

不过,当计算图中出现菱形依赖时,要是直接在backward中反向传播就会导致指数级增长的运算量。考虑以下计算图:

  1. val weight: DoubleLayer = DoubleWeight(math.random)
  2. def diamondComputationalGraph: DoubleLayer = {
  3. val square0: DoubleLayer = weight * weight
  4. val square1: DoubleLayer = square0 * square0
  5. val square2: DoubleLayer = square1 * square1
  6. val square3: DoubleLayer = square2 * square2
  7. square3
  8. }

当调用diamondComputationalGraph.train训练时,在反向传播阶段,首先触发square3backward,然后square3会调用左右操作数的backward进行反向传播。由于square3的左右操作数都是square2,所以square2backward会触发两次。而每次square2backward会触发两次square1backwardsquare1backward一共就触发了4次。以此类推,最终weightbackward会触发16次。

方案三:基于引用计数的Resource

为了解决这个问题,我们增加了一个引用计数功能,我们让forward返回的不是Tape,而是带有引用计数的Resource

  1. case class Resource[A](value: A, release: UnitContinuation[Unit])

那么现在的乘法运算符可以实现成这样:

  1. def *(operand0: DoubleLayer, operand1: DoubleLayer): DoubleLayer = {
  2. var referenceCount = 0
  3. var deltaAccumulator = 0.0
  4. def forward: UnitContinuation[Resource[Tape[Double, Double]]] = monadic[UnitContinuation] {
  5. referenceCount += 1
  6. val resource0 = operand0.forward.each
  7. val resource1 = operand1.forward.each
  8. val data = resource0.value.data * resource1.value.data
  9. val backward = { delta =>
  10. monadic[UnitContinuation] {
  11. deltaAccumulator += delta.each
  12. }
  13. }
  14. val tape = Tape(data, backward)
  15. val release = monadic[UnitContinuation] {
  16. referenceCount -= 1
  17. if (referenceCount == 0) {
  18. resource0.value.backward(deltaAccumulator * resource1.value.data)
  19. resource1.value.backward(deltaAccumulator * resource0.value.data)
  20. deltaAccumulator = 0.0
  21. resource0.release().each
  22. resource1.release().each
  23. }
  24. }
  25. Resource(tape, release)
  26. }
  27. DoubleLayer(forward)
  28. }

相比第二版的乘法,我们做了如下修改:
* 每个DoubleLayer会记录forward的次数。
* backward不再立即触发递归的反向传播,而是只把delta暂存到deltaAccumulator
* 增加了release函数,release函数会释放引用计数,当引用计数释放到零时,触发反向传播并且释放上游的引用计数。

train也需要修改,在发起反向传播之处时,得在backward之后调用一次release

  1. def train(lossFunction: DoubleLayer): UnitContinuation[Data] = monadic[UnitContinuation] {
  2. val resource: Resource[Tape[Data, Delta]] = lossFunction.forward.each
  3. val data = resource.data
  4. resource.value.backward(1.0).each
  5. resource.release().each
  6. }

经过这样的修改,现在diamondComputationalGraph.train训练时,只调用了一次lossFunction.forward,所以resource的引用计数为1。

然后,当train调用resource.value.backward时,新的backward实现只会把delta暂存到deltaAccumulator,不触发递归的反向传播。

最后,当train调用release时,引用计数减为零,触发递归的反向传播。

当反向传播到square3时,由于square2既是square3的左操作数也是右操作数,所以square2会被square3引用两次,引用计数为2。当square3触发release时,会同时触发square3的左操作数和右操作数的release,也就导致square2两遍,刚好可以把引用计数减为零,继续触发square2的反向传播。

以此类推,计算图的每个节点都有引用计数。即使在菱形计算图的情况下,也可以保证只在最后一次release时才触发反向传播,而不会导致指数增长的运算量。

问题:动态神经网络的引用计数

这种基于引用计数的算法,必须保证forward调用次数和release相同。然而,在复杂的动态网络中很难保证。比如:

  1. val weight: DoubleLayer = DoubleWeight(math.random)
  2. def diamondComputationalGraph: DoubleLayer = DoubleLayer(monadic[UnitContinuation] {
  3. if (weight.forward.each.value.data > 0.5) {
  4. val square0: DoubleLayer = weight * weight
  5. if (square0.forward.each.value.data > 0.5) {
  6. val square1: DoubleLayer = square0 * square0
  7. if (square1.forward.each.value.data > 0.5) {
  8. square1 * square1
  9. } else {
  10. square1
  11. }
  12. } else {
  13. square0
  14. }
  15. } else {
  16. weight
  17. }
  18. })

以上动态神经网络中调用了3次forward,但却无处可以调用对应的release。这样写法会导致引用计数没有机会清零,从而导致无法反向传播。

方案四:自动管理引用计数的ResourceT

我们实现了一套ResourceT Monad Transformer,可以用来自动生成对应与forwardrelease

在这一方案中,我们把DoubleLayer.forward的类型改为了ResourceT[UnitContinuation, Tape[Double, Double]]

  1. trait DeepLearning[Differentiable] {
  2. type Data
  3. type Delta
  4. def forward(differentiable: Differentiable): ResourceT[UnitContinuation, Tape[Data, Delta]]
  5. }

ResourceT是个opacity type alias,底层实际类型是UnitContinuation[Resource[Tape[Double, Double]]]

UnitContinuation是DeepLearning.scala使用的异步编程库。加入UnitContinuation还可以解决forward次数统计的问题。因为UnitContinuation是个惰性计算的Monad,要等到实际运行前才会计算forward次数,这样一来,在动态神经网络中,没有经过的分支就没有开销,运算量小了很多。

由于ResourceT是个Monad,所以我们可以把动态神经网络改成这样:

  1. val weight: DoubleLayer = DoubleWeight(math.random)
  2. def diamondComputationalGraph: DoubleLayer = {
  3. val resourcet: ResourceT[UnitContinuation, Tape[Double, Double]] = weight.forward
  4. resourcet.flatMap { tape =>
  5. if (tape.data > 0.5) {
  6. val square0: DoubleLayer = weight * weight
  7. square0.forward.flatMap { tape0 =>
  8. if (tape0.data > 0.5) {
  9. val square1: DoubleLayer = square0 * square0
  10. square1.forward.flatMap { tape1 =>
  11. if (tape1.data > 0.5) {
  12. (square1 * square1).forward
  13. } else {
  14. square1.forward
  15. }
  16. }
  17. } else {
  18. square0.forward
  19. }
  20. }
  21. } else {
  22. weight.forward
  23. }
  24. }
  25. }

神经网络与函数式编程(四)Monadic Deep Learning中,我们介绍过flatMapMonad.bind实现,作用是生成一个两阶段计算图。

具体的resourcet.flatMap会转发到Monad[ResourceT[UnitContinuation, ?]].bind

ResourceT作为一个Monad Transformer,其bind实现会额外增加一个功能,把release收集起来:

  1. def bind[A, B](rfa: ResourceT[F, A])(f: (A) => ResourceT[F, B]): ResourceT[F, B] = {
  2. val ResourceT(fa) = rfa
  3. fa.flatMap { resourceA =>
  4. val ResourceT(fb) = f(resourceA.value)
  5. fb.map { resourceB =>
  6. val release = resourceB.release >> resourceA.release
  7. Resource(resourceB.value, release)
  8. }
  9. }
  10. ResourceT(result)
  11. }

这样一来,当用户调用resourcet.flatMap时,返回的ResourceT所包含的release操作由两阶段的release组成:

  1. val release = resourceB.release >> resourceA.release

通过引入ResourceT来管理Resource,我们把release的引用关系自动记录下来了。在每次forward.flatMap时记录引用,最后,在整个网络train时再触发release,所以ResourceT的用户不需要手动调用release,解决了动态神经网络的引用计数问题。

方案五:增加异常处理的Monad Transformer

DeepLearning.scala除了需要引用计数之外,还需要异常处理。和ResourceT类似,这也可以通过增加一层TryT Monad Transformer实现。

这样一来,forward返回的类型变得更加复杂了。

  1. trait DeepLearning[Differentiable] {
  2. type Data
  3. type Delta
  4. def forward(differentiable: Differentiable): TryT[ResourceT[UnitContinuation, ?], Tape[Data, Delta]]
  5. }

方案六:简化Monad类型

在加入支持异步、引用计数管理之后,我们的计算图类型已经非常复杂,变成了TryT[ResourceT[UnitContinuation, ?], Tape[Data, Delta]],这回给用户带来不必要的学习负担。所以我们创建了一个Do类型。

Do是个opacity type alias。对于任何Do[A],底层实际类型是TryT[ResourceT[UnitContinuation, ?], A]

然后我们可以把所有用到TryTResourceT之处,改为使用Do

  1. trait DeepLearning[Differentiable] {
  2. type Data
  3. type Delta
  4. def forward(differentiable: Differentiable): Do[Tape[Data, Delta]
  5. }

方案六是DeepLearning.scala 2.0最终的计算图表示方式。

结论

在DeepLearning.scala 2.0中,我们用Monad表示计算图,把每一个计算图上的特性实现为一个Monad或者一个Monad Transformer,比如:

然后用Do把这些Monad Transformer封装起来。最终,DeepLearning.scala的用户并不需要知道Do的内部构造,就能够方便的使用Do构造出动态计算图,自动获得上述所有功能。

在本篇文章中,我们揭示了在Do的简单用法背后隐藏的层层Monad。在下一篇文章中,我将揭示DeepLearning.scala另一个功能,插件系统,背后隐藏的底层机制。

添加新批注
在作者公开此批注前,只有你和作者可见。
回复批注