![ba4239886e51fe14f69f5e461cdb7c26.png](https://i-blog.csdnimg.cn/blog_migrate/00679b088dd1f36da1f5fa9c85eca1fa.jpeg)
这个章节是微分编程的最后一个章节,前面两个章节细数了微分编程的“三宗罪”,可逆编程,可逆编程那部分把很多干巴巴的理论甚至硬件都扯完了,所以今天就可以很轻松的从例子出发讲讲几种AD机制的不一样的地方,为什么会有这么大的性能差异。今天会涉及的几种AD机制都是源到源的AD,包括基于checkpointing的adjoint mode AD,基于可逆编程的adjoint mode AD和tangent mode AD三种。
相信很多人都知道第一类贝塞尔函数,它在物理里面是重要的函数。一个
其中Gamma函数
性能测试总结表
![d751ebd40fc5114bb692ed77a70aa9e5.png](https://i-blog.csdnimg.cn/blog_migrate/450f89203939c82382b0ba35d7dae84c.jpeg)
其中第一行Julia的不带AD的forward实现是基准,NiLang是可逆编程嵌入式语言(eDSL),NiLang.AD是该eDSL上的Adjoint mode AD的实现。ForwardDiff是Julia下Tangent mode AD的实现。Zygote是Julia下源到源的Adjoint mode AD的实现。Manual是手动绑定(或者说operator overloading)的Adjoint mode AD,定义见文中。
正文
- 基于计算图的AD:Zygote
上述贝塞尔函数的表达式是个求和,每个求和项可以表达为
function besselj(ν, z; atol=1e-8)
k = 0
s = (z/2)^ν / factorial(ν)
out = s
while abs(s) > atol
k += 1
s *= (-1) / k / (k+ν) * (z/2)^2
out += s
end
out
end
如果调用了自动微分库Julia的源到源自动微分库Zygote,它会劫持了这段代码的SSA表达式,并把它表达成了计算图。这个函数的计算图被画在了图一,它把数据转换成了箭头,函数转换成了节点,并展开了流控制令计算图变得很长。它没有内存更没有inplace函数的概念,像极了数学表达式,所以很容易通过链式法则来求导。
![ae090af215903238790ef3003060183a.png](https://i-blog.csdnimg.cn/blog_migrate/bf36fbb33305161aa499d037c49020c1.jpeg)
不费吹灰之力,就可以通过如下短短一行代码得到第一类贝塞尔函数的导数
julia> using Zygote
julia> Zygote.gradient(besselj, 2, 1.0)
(-0.07635849493801533, 0.21024361585183118)
源到源自动微分真是太有趣了,它继tensorflow,pytorch这些基于operator overloading的机制之后,让自动微分变得更加自动了,大大减少获得导数的劳动力成本。同样的源到源自动微分还有Tapenade等。正当我们感慨源到源的美妙时,我们还关心的问题是,它性能可以吗?于是我们做了个数值小实验,
julia> using BenchmarkTools
julia> @benchmark besselj(2, 1.0)
BenchmarkTools.Trial:
memory estimate: 0 bytes
allocs estimate: 0
--------------
minimum time: 23.363 ns (0.00% GC)
median time: 23.731 ns (0.00% GC)
mean time: 26.207 ns (0.00% GC)
maximum time: 89.767 ns (0.00% GC)
--------------
samples: 10000
evals/sample: 996
julia> @benchmark gradient(besselj, 2, 1.0)
BenchmarkTools.Trial:
memory estimate: 13.47 KiB
allocs estimate: 411
--------------
minimum time: 37.788 μs (0.00% GC)
median time: 38.771 μs (0.00% GC)
mean time: 44.812 μs (5.75% GC)
maximum time: 9.084 ms (98.76% GC)
--------------
samples: 10000
evals/sample: 1
end{
lstlisting}
end{
minipage}
我们发现Zygote比前向计算慢了1500倍以上,不仅如此,还多了13KB的heap上的内存分配,而前向过程根本没有heap上的内存分配。造成这个结果的原因有两部分,一部分是Zygote过于依赖Julia的Compiler去推导类型并去除运行时的overhead,但是Julia的type inference并没有做到完美。另一个比较本质的原因是,这样的自动微分用了naive的checkpointing,造成了很多内存分配,虽然自动了,但是性