如果你有使用过诸如PyTorch,Flux/Tracker,AutoGrad这一类基于运算符重载的自动微分库,就会发现,这些库有两个通病:
只能使用框架所提供的函数和矩阵/张量类型,如果想要对一般的程序进行求导就不行了
无法处理控制流,因为简单的运算符重载无法记录下来控制流
在Yan LeCun等人的号召下,想要实现可微分编程(Differentiable Programming)如果没有上面这两个功能可不行。不能对控制流进行微分叫什么可微分编程?因为我们希望我们编写的任意在数学上成立的程序都可以进行自动微分。此外上篇文章的评论区里有人提到了TensorFlow的自动产生符号导数然后加入计算图的功能。而实际上如果对编译器数学的同学会知道这类符号计算就是一个简单的编译步骤。
这些问题在源对源(source to source)的自动微分下实现起来将非常自然,这篇文章将实现一个不带控制流的简单版本,而完整的版本已经在Julia中通过staged programming的方式实现了,这个包叫做 Zygote.jl 。它中文翻译很搞笑,叫卵子。使用这个包的人都用了个卵子自动微分。我在今年JuliaCon的周五Hackthon期间在Zygote的作者Mike的帮助下实现了一个简单版本的。
不过,在我们开始编写程序前,让我们来回顾一些基础知识。
Julia语言的编译过程
首先让我们来简单了解一下Julia语言是如何进行编译的。
首先,所有的代码本质上都是一些字符串(string),存储在硬盘上的文本文件中
我们首先要解析(parse)这些字符串,得到一个抽象语法数(Abstract Syntax Tree,AST)
而 AST 里有一些节点是宏,这些宏是一些只接受编译时期变量的,里面描述了如何产生更多的 AST,在这一步将会运行这些宏,我们成为 展开AST。你可以通过 macroexpand 宏查看这一步的编译结果
这个时候我们再将AST里的语法糖等节点全部替换为函数调用,并且使用SSA(Static Single Assignment)形式的IR作为更低级的表示。什么是SSA IR?我们将在后面介绍
到此位置我们完成了代码的初始化过程。
然后我们的函数会在被派发的时候才会被继续编译,这是因为对于一般的函数(generic function)我们是无法在编译时期就确定这个函数的变量类型的,从而无法产生定制的机器