API统一、干净,新型EagerPy实现多框架无缝衔接

近年来,深度学习领域的进展与深度学习框架的开发同步进行 。这些框架为自动微分和 GPU 加速提供了高级且高效的 API,从而可以利用相对较少和简单的代码实现极度复杂和强大的深度学习模型 。
【API统一、干净,新型EagerPy实现多框架无缝衔接】最初,Theano、Caffe、MXNet、TensorFlow 和 CNTK 等很多流行的深度学习框架使用的是基于图的方法 。用户首先需要定义一个静态数据流图(static data flow graph),然后可以对它进行高效地微分、编译并在 GPU 上执行 。所以,提前了解整个计算图有助于实现高性能 。
但是,这种方法导致难以调试模型以及实现具有变化图(changing graph)的动态模型(如 RNN) 。
所以,针对这种方法的局限性,深度学习模型的 Eager Execution 成为了深度学习研究领域的主流方法 。用户不再需要提前构建静态数据流图,Eager Execution 框架自身就可以提供 define-by-run 的 API,它可以高速地构建临时的动态图 。目前,两大主流深度学习框架 PyTorch 和 TensorFlow 都在使用 eager execution 方法 。
而在本文中,来自德国图宾根大学和图宾根伯恩斯坦计算神经科学中心的研究者将 eager execution 进行了扩展,提供了一个新的 Python 框架 EagerPy,它可以编写自动且原生地适配 PyTorch、TensorFlow、Jax 和 Numpy 的代码 。EagerPy 对库开发者和用户都有裨益 。

API统一、干净,新型EagerPy实现多框架无缝衔接

文章插图
 
论文地址:https://arxiv.org/abs/2008.04175v1
项目地址:https://github.com/jonasrauber/eagerpy
EagerPy 能够编写与框架无关(framework-agnostic)的代码,这些代码可以与 PyTorch、TensorFlow、Jax 和 NumPy 实现原生地适配 。
这样一来,首先对于新库开发者而言,他们不仅可以选择同时支持上述这几个主流深度学习框架或者为每个框架重新实现库,而且可以对代码重复进行处理 。
其次对于这些库的使用者而言,他们也可以更轻松地切换深度学习框架,并且不会被特定的第三方库锁定 。
不仅如此,单个框架的使用者也会从 EagerPy 中获益,这是因为 EagerPy 提供了全面的类型注释以及对方法链接到任何框架的一致支持 。
接下来我们来看 EagerPy 的具体设计与实现 。
EagerPy 的设计与实现
EagerPy 的构建考虑到了 4 个设计目标 。两个主要的目标是为需要执行操作的人提供统一的 API,并维护底层框架的原始性能 。这两个主要目标定义了 EagerPy 是什么,所以是设计的核心 。
与底层框架特定的 API 相比,完全可链接的 API 和全面的类型检查支持这两个附加目标使 EagerPy 更加易于使用,也更安全 。
尽管进行了这些更改和改进,但研究者尝试避免不必要的熟悉度(familiarity)损失 。只要有意义,EagerPy API 都会遵循 NumPy、PyTorch 和 JAX 设置的标准 。
统一的 API
为了实现语法上的一致性,研究者使用适当的方法定义了一个抽象 Tensor 类,并使用一个实例变量来保存原生张量(native tensor),然后为每个支持的框架实现一个特定的子类 。对于诸如 sum 或 log 的很多操作,这就像调用底层框架一样简单;而对于其他操作,则工作量会稍大一些 。
最困难的部分是统一自动微分 API 。PyTorch 使用了一个低级的 autograd API,该 API 允许但也需要对反向传播的精确控制 。TensorFlow 使用基于梯度磁带(gradient tapes)的更高级 API 。而 JAX 使用基于微分函数的相当高级的 API 。
所以,为了统一它们,EagerPy 模仿了 JAX 的高级功能 API,并在 PyTorch 和 TensorFlow 中重新实现 。EagerPy 通过 value_and_grad_fn() 函数将其开放 。
此外,能够编写自动与所有支持的框架一起运行的代码,不仅需要语法,还需要语义统一 。为了确保这一点,EagerPy 附带了一个庞大的测试套件,该套件可以验证不同框架特定子类之间的一致性 。它会在所有 pull-request 上自动运行,并且需要通过之后才能合并新代码 。
测试套件还可以作为所支持的操作和参数组合的最终参考 。这样就可以避免文档和实现之间出现不一致,并在实践中引出测试驱动开发过程 。
原始性能
没有 EagerPy,想要与不同深度学习框架进行交互的代码必须经过 NumPy 实现 。这需要在 CPU(NumPy)和 GPU(PyTorch、TensorFlow 和 JAX)之间进行高成本的内存复制,反之亦然 。
此外,许多计算仅在 CPU 上执行,为了避免这种情况,EagerPy 仅保留对原始框架特定张量的引用(例如 GPU 上的 PyTorch 张量),并将所有的操作委托给相应的框架 。这几乎不产生任何的计算开销 。


推荐阅读