@atry
2017-11-16T10:06:20.000000Z
字数 9875
阅读 1339
神经网络与函数式编程
在本系列先前的文章神经网络与函数式编程(四)Monadic Deep Learning中,我们学习了DeepLearning.scala如何实现多态函数,把各种不同的参数统一转为Do[Tape]
来处理。在本篇文章中,我们将讨论DeepLearning.scala的反向传播机制,看看DeepLearning.scala如何实现Monadic风格的可差分计算图。
Tape
的反向传播首先我将简化一下DeepLearning.scala的实现,把涉及Monad的部分去掉,展示一下如何DeepLearning.scala内部的反向传播流程。
以一元线性回归模型为例,如果采用LAD惩罚函数,可以用DeepLearning.scala表达成这样:
val weight = DoubleWeight(math.random)
def f(x: Double): DoubleLayer = x * weight
def lossFunction(x: Double, expectedY: Double): DoubleLayer = {
abs(expectedY - f(x))
}
给定带标记的数据样本x
和expectedY
,可以用train
函数训练:
val x: Double = ???
val expectedY: Double = ???
lossFunction(x, expectedY).train
其中,lossFunction(x, expectedY)
是一个计算图DoubleLayer
。在本系列前一篇文章,我们已经知道DoubleLayer
可以forward
转为Do[Tape[Double, Double]]
。Do
是个Monad
类型,我们暂时忽略,把DoubleLayer.forward
简单的看成Tape[Double, Double]
。简化后的Tape
定义如下:
case class Tape[Data, Delta](data: Data, backward: Delta => Unit)
Tape
包含前向阶段的结果data
,以及用来进行反向传播的backward
函数。
比如train
是反向传播的起点,可以实现成这样:
trait DeepLearning[Differentiable] {
type Data
type Delta
def forward(differentiable: Differentiable): Tape[Data, Delta]
def train(lossFunction: DoubleLayer): Data = {
val tape: Tape[Data, Delta] = lossFunction.forward
tape.backward(1.0)
tape.data
}
}
train
用到的lossFunction
展开后是abs(expectedY - x * weight)
。
// TODO: 此处缺一张图
其中可差分的abs
实现如下
def abs(operand0: DoubleLayer): DoubleLayer = {
def forward: Tape[Double, Double] = {
val tape = operand0.forward
val data = math.abs(tape.data)
val backward = { delta =>
tape.backward(math.sign(tape.data) * delta)
}
Tape(data, backward)
}
DoubleLayer(forward)
}
可以看到在前向阶段,abs
的data
是operand0
的data的绝对值。而在反向传播时,delta
函数把导数经过计算后,传给operand0
的backward
。
而abs
用到的operand0
则是可差分的减法,根据求导公式可以写成这样:
def -(operand0: DoubleLayer, operand1: DoubleLayer): DoubleLayer = {
def forward: Tape[Double, Double] = {
val tape0 = operand0.forward
val tape1 = operand1.forward
val data = tape0.data - tape1.data
val backward = { delta =>
tape0.backward(delta)
tape1.backward(-delta)
}
Tape(data, backward)
}
DoubleLayer(forward)
}
注意,减法与abs
不同,是个二元运算符,所以backward
时需要分别反向传播到两个操作数的backward
。
然后,减法的operand1
是f(x)
,是个可差分的乘法:
def *(operand0: DoubleLayer, operand1: DoubleLayer): DoubleLayer = {
def forward: Tape[Double, Double] = {
val tape0 = operand0.forward
val tape1 = operand1.forward
val data = tape0.data * tape1.data
val backward = { delta =>
tape0.backward(delta * tape1.data)
tape1.backward(delta * tape0.data)
}
Tape(data, backward)
}
DoubleLayer(forward)
}
最后,weight
是个DoubleWeight
。DoubleWeight
的backward
会修改自己内部保存的值:
trait DoubleWeight {
var data: Double
def forward: Tape[Double, Double] = {
val backward = { delta =>
data -= delta
}
Tape(data, backward)
}
}
到此为止,可以看到,当用户调用train
时。
train
内部首先会递归调用forward
,在最深处由DoubleWeight
取得data
,然后一层一层执行前向阶段,返回一个Tape
,其中包含前些阶段的执行结果。
然后,train
内部会递归调用Tape
上的backward
,把导数反向传播,直到传到DoubleWeight
上,令DoubleWeight
得到训练。
到目前为止,我们实现的计算图是同步执行的。这意味着,计算图上的节点都不可能并行执行,即使这些节点之间并没有相互依赖关系。为了实现神经网络与函数式编程(四)Monadic Deep Learning中提到的粗粒度并行计算,我们需要让forward
返回一个异步任务,而不应该阻塞。
trait DeepLearning[Differentiable] {
type Data
type Delta
def forward(differentiable: Differentiable): UnitContinuation[Tape[Data, Delta]]
def train(lossFunction: DoubleLayer): UnitContinuation[Data] = monadic[UnitContinuation] {
val tape = lossFunction.forward.each
tape.backward(UnitContinuation.now(1.0)).each
tape.data
}
}
以上代码中,我们引入了UnitContinuation来表示异步任务。train
也做了修改,改用monadic风格的控制流。
相应的Tape.backward
也需要改为异步结构。
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风格,比如乘法操作:
def *(operand0: DoubleLayer, operand1: DoubleLayer): DoubleLayer = {
def forward: UnitContinuation[Tape[Double, Double]] = {
val Parallel(result) = (Parallel(operand0.forward) |@| Parallel(operand1.forward)) { (tape0, tape1) =>
val data = tape0.data * tape1.data
val backward = { deltaContinuation =>
deltaContinuation.flatMap { delta =>
val backward0 = tape0.backward(delta * tape1.data)
val backward1 = tape1.backward(delta * tape0.data)
val Parallel(backwardBoth) = Parallel(backward0) |@| Parallel(backward1) { (_, _) => }
backwardBoth
}
}
Tape(data, backward)
}
result
}
DoubleLayer(forward)
}
注意,我们实现了乘法这样的二元操作符时,无论是前向阶段还是反向阶段,都利用Parallel
让两个操作数上的计算得到并行执行。
到此为止,方案二通过引入UnitContinuation
,让所有计算图都得到了异步计算的能力,并且可以并行执行多个计算图的分支。
与autograd等早期的自动求导算法不同,DeepLearning.scala在二元操作处,能对两个操作数都求导。你可以在减法和乘法的上看到,backward
内部会分别触发两个操作数的backward
。这种做法可以支持复杂的网络结构.
不过,当计算图中出现菱形依赖时,要是直接在backward
中反向传播就会导致指数级增长的运算量。考虑以下计算图:
val weight: DoubleLayer = DoubleWeight(math.random)
def diamondComputationalGraph: DoubleLayer = {
val square0: DoubleLayer = weight * weight
val square1: DoubleLayer = square0 * square0
val square2: DoubleLayer = square1 * square1
val square3: DoubleLayer = square2 * square2
square3
}
当调用diamondComputationalGraph.train
训练时,在反向传播阶段,首先触发square3
的backward
,然后square3
会调用左右操作数的backward
进行反向传播。由于square3
的左右操作数都是square2
,所以square2
的backward
会触发两次。而每次square2
的backward
会触发两次square1
的backward
。square1
的backward
一共就触发了4次。以此类推,最终weight
的backward
会触发16次。
Resource
为了解决这个问题,我们增加了一个引用计数功能,我们让forward
返回的不是Tape
,而是带有引用计数的Resource
:
case class Resource[A](value: A, release: UnitContinuation[Unit])
那么现在的乘法运算符可以实现成这样:
def *(operand0: DoubleLayer, operand1: DoubleLayer): DoubleLayer = {
var referenceCount = 0
var deltaAccumulator = 0.0
def forward: UnitContinuation[Resource[Tape[Double, Double]]] = monadic[UnitContinuation] {
referenceCount += 1
val resource0 = operand0.forward.each
val resource1 = operand1.forward.each
val data = resource0.value.data * resource1.value.data
val backward = { delta =>
monadic[UnitContinuation] {
deltaAccumulator += delta.each
}
}
val tape = Tape(data, backward)
val release = monadic[UnitContinuation] {
referenceCount -= 1
if (referenceCount == 0) {
resource0.value.backward(deltaAccumulator * resource1.value.data)
resource1.value.backward(deltaAccumulator * resource0.value.data)
deltaAccumulator = 0.0
resource0.release().each
resource1.release().each
}
}
Resource(tape, release)
}
DoubleLayer(forward)
}
相比第二版的乘法,我们做了如下修改:
* 每个DoubleLayer
会记录forward
的次数。
* backward
不再立即触发递归的反向传播,而是只把delta
暂存到deltaAccumulator
。
* 增加了release
函数,release
函数会释放引用计数,当引用计数释放到零时,触发反向传播并且释放上游的引用计数。
train
也需要修改,在发起反向传播之处时,得在backward
之后调用一次release
:
def train(lossFunction: DoubleLayer): UnitContinuation[Data] = monadic[UnitContinuation] {
val resource: Resource[Tape[Data, Delta]] = lossFunction.forward.each
val data = resource.data
resource.value.backward(1.0).each
resource.release().each
}
经过这样的修改,现在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
相同。然而,在复杂的动态网络中很难保证。比如:
val weight: DoubleLayer = DoubleWeight(math.random)
def diamondComputationalGraph: DoubleLayer = DoubleLayer(monadic[UnitContinuation] {
if (weight.forward.each.value.data > 0.5) {
val square0: DoubleLayer = weight * weight
if (square0.forward.each.value.data > 0.5) {
val square1: DoubleLayer = square0 * square0
if (square1.forward.each.value.data > 0.5) {
square1 * square1
} else {
square1
}
} else {
square0
}
} else {
weight
}
})
以上动态神经网络中调用了3次forward
,但却无处可以调用对应的release
。这样写法会导致引用计数没有机会清零,从而导致无法反向传播。
ResourceT
我们实现了一套ResourceT
Monad Transformer,可以用来自动生成对应与forward
的release
。
在这一方案中,我们把DoubleLayer.forward
的类型改为了ResourceT[UnitContinuation, Tape[Double, Double]]
。
trait DeepLearning[Differentiable] {
type Data
type Delta
def forward(differentiable: Differentiable): ResourceT[UnitContinuation, Tape[Data, Delta]]
}
ResourceT是个opacity type alias,底层实际类型是UnitContinuation[Resource[Tape[Double, Double]]]
。
UnitContinuation是DeepLearning.scala使用的异步编程库。加入UnitContinuation
还可以解决forward
次数统计的问题。因为UnitContinuation
是个惰性计算的Monad
,要等到实际运行前才会计算forward
次数,这样一来,在动态神经网络中,没有经过的分支就没有开销,运算量小了很多。
由于ResourceT
是个Monad
,所以我们可以把动态神经网络改成这样:
val weight: DoubleLayer = DoubleWeight(math.random)
def diamondComputationalGraph: DoubleLayer = {
val resourcet: ResourceT[UnitContinuation, Tape[Double, Double]] = weight.forward
resourcet.flatMap { tape =>
if (tape.data > 0.5) {
val square0: DoubleLayer = weight * weight
square0.forward.flatMap { tape0 =>
if (tape0.data > 0.5) {
val square1: DoubleLayer = square0 * square0
square1.forward.flatMap { tape1 =>
if (tape1.data > 0.5) {
(square1 * square1).forward
} else {
square1.forward
}
}
} else {
square0.forward
}
}
} else {
weight.forward
}
}
}
在神经网络与函数式编程(四)Monadic Deep Learning中,我们介绍过flatMap
由Monad.bind
实现,作用是生成一个两阶段计算图。
具体的resourcet.flatMap
会转发到Monad[ResourceT[UnitContinuation, ?]].bind
。
ResourceT
作为一个Monad Transformer,其bind
实现会额外增加一个功能,把release
收集起来:
def bind[A, B](rfa: ResourceT[F, A])(f: (A) => ResourceT[F, B]): ResourceT[F, B] = {
val ResourceT(fa) = rfa
fa.flatMap { resourceA =>
val ResourceT(fb) = f(resourceA.value)
fb.map { resourceB =>
val release = resourceB.release >> resourceA.release
Resource(resourceB.value, release)
}
}
ResourceT(result)
}
这样一来,当用户调用resourcet.flatMap
时,返回的ResourceT
所包含的release
操作由两阶段的release
组成:
val release = resourceB.release >> resourceA.release
通过引入ResourceT
来管理Resource
,我们把release
的引用关系自动记录下来了。在每次forward.flatMap
时记录引用,最后,在整个网络train
时再触发release
,所以ResourceT
的用户不需要手动调用release
,解决了动态神经网络的引用计数问题。
DeepLearning.scala除了需要引用计数之外,还需要异常处理。和ResourceT
类似,这也可以通过增加一层TryT Monad Transformer实现。
这样一来,forward
返回的类型变得更加复杂了。
trait DeepLearning[Differentiable] {
type Data
type Delta
def forward(differentiable: Differentiable): TryT[ResourceT[UnitContinuation, ?], Tape[Data, Delta]]
}
在加入支持异步、引用计数管理之后,我们的计算图类型已经非常复杂,变成了TryT[ResourceT[UnitContinuation, ?], Tape[Data, Delta]]
,这回给用户带来不必要的学习负担。所以我们创建了一个Do类型。
Do
是个opacity type alias。对于任何Do[A]
,底层实际类型是TryT[ResourceT[UnitContinuation, ?], A]
。
然后我们可以把所有用到TryT
、ResourceT
之处,改为使用Do
:
trait DeepLearning[Differentiable] {
type Data
type Delta
def forward(differentiable: Differentiable): Do[Tape[Data, Delta]
}
方案六是DeepLearning.scala 2.0最终的计算图表示方式。
在DeepLearning.scala 2.0中,我们用Monad表示计算图,把每一个计算图上的特性实现为一个Monad或者一个Monad Transformer,比如:
UnitContinuation
ResourceT
TryT
然后用Do
把这些Monad Transformer封装起来。最终,DeepLearning.scala的用户并不需要知道Do
的内部构造,就能够方便的使用Do
构造出动态计算图,自动获得上述所有功能。
在本篇文章中,我们揭示了在Do
的简单用法背后隐藏的层层Monad
。在下一篇文章中,我将揭示DeepLearning.scala另一个功能,插件系统,背后隐藏的底层机制。