杜伟、小舟编辑

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

这个新型 Python 框架对库开发者和用户都大有裨益。

近年来,深度学习领域的进展与深度学习框架的开发同步进行。这些框架为自动微分和 GPU 加速提供了高级且高效的 API,从而可以利用相对较少和简单的代码实现极度复杂和强大的深度学习模型。

最初,Theano、Caffe、MXNetTensorFlow 和 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 对库开发者和用户都有裨益。

  • 论文地址: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 张量),并将所有的操作委托给相应的框架。这几乎不产生任何的计算开销。

完全可链接的 API

求和或平方之类的许多运算都要采用张量并返回一个张量。通常情况下,这些运算按顺序被调用。例如使用平方、求和和开平方根以计算 L2 范数

在 EagerPy 中,所有运算都成为了张量对象(tensor object)上可用的方法。这样就可以按照它们的自然顺序(x.square().sum().sqrt())来链接操作。相反,例如,NumPy 需要相反的操作顺序,即 np.sqrt(np.square(x).sum())。

类型检查

在 Python3.5 中,Python 语法的扩展已经实现了对类型注释的支持(van Rossum 等人,2015 年)。即使具有类型注释,Python 仍然是一种动态类型化的编程语言,并且当前在运行时会忽略所有类型注释。但是,我们可以在运行代码之前通过静态代码分析器检查这些类型注释。

EagerPy 带有所有参数和返回值的全面类型注释,并使用 Mypy(Lehtosalo 等人,2016 年)对这些注释进行检查。这有助于我们捕获 EagerPy 中的漏洞,否则这些漏洞将一直不会被发现。

EagerPy 用户可以通过键入自己代码的注释,并根据 EagerPy 的函数签名(function signature)自动检查代码来进一步优化。这一点很关键,因为 TensorFlow、NumPy 和 JAX 当前自身不提供类型注释。

EagerPy 的代码实例解析

如下代码 1 为一个通用 EagerPy 范数函数,它可以通过任何框架中的原生张量被调用,并且返回的范数依然作为同一个框架中的原生张量

代码 1:框架无关的范数函数。

EagerPy 和原生张量之间的转换

原生张量可以是 PyTorch GPU 或 CPU 张量,如下代码 2 所示:

代码 2:原生 PyTorch 张量

可以是 TensorFlow 张量,如下代码 3 所示:

代码 3:原生 TensorFlow 张量


可以是 JAX 数组,如下代码 4 所示:

代码 4:原生 JAX 数组。

可以是 NumPy 数组,如下代码 5 所示:

代码 5:原生 NumPy 数组。

无论是哪种原生张量,通常都可以使用 ep.astensor 将它转换为适当的 EagerPy 张量。在此步骤中,通过使用正确的 EagerPy 张量类来自动封装原生张量。此外,最初的原生张量通常可以利用. raw 属性实现访问。完整示例如下代码 6 所示:

EagerPy 和原生张量之间的转换。

在函数中通常将所有输入转换为 EagerPy 张量。这可以通过单独调用 ep.astensor 完成,但在使用 ep.astensors 时,代码可以更加简洁,如下:

实现框架无关的通用函数

通过上文中的转换函数,我们可以定义一个简单的框架无关函数,如下代码 8 所示:

代码 8:一个简单的框架无关范数函数。

如下代码 9 所示,通过一个 PyTorch 张量来调用范数函数:

 如下代码 10 所示,通过一个 TensorFlow 张量来调用范数函数:

此外,还需要注意一点,如果如上代码 8 所示使用 EagerPy 张量来调用函数,则 ep.astensor 调用只会返回它的输入。但是,最后一行代码中的 result.raw 调用依然会提取底层原生张量。通常而言,实现的通用函数最好可以透明地操控任何原生张量和 EagerPy 张量,也就是说返回类型应该总是与输入类型相匹配。

这在 Foolbox 等库中非常有用,可以使用户同时处理 EagerPy 和原生张量

为此,EagerPy 提供上述转换函数的两种派生函数,分别是 ep.astensor_和 ep.astensors_,它们可以返回一个能够恢复输入类型的反转函数。

如果 astensor_的输入是一个原生张量,则 restore_type 等同于. raw;而如果原输入是一个 EagerPy 张量,则 restore_type 将不会调用. raw。因此,我们可以编写对任何输入都透明的改进版框架无关通用函数,如下代码 11 所示:

最后,如下代码 12 所示,使用 ep.astensors_来转换和恢复多个输入:


理论框架EagerPyPython
相关数据
深度学习技术

深度学习(deep learning)是机器学习的分支,是一种试图使用包含复杂结构或由多重非线性变换构成的多个处理层对数据进行高层抽象的算法。 深度学习是机器学习中一种基于对数据进行表征学习的算法,至今已有数种深度学习框架,如卷积神经网络和深度置信网络和递归神经网络等已被应用在计算机视觉、语音识别、自然语言处理、音频识别与生物信息学等领域并获取了极好的效果。

范数技术

范数(norm),是具有“长度”概念的函数。在线性代数、泛函分析及相关的数学领域,是一个函数,其为向量空间内的所有向量赋予非零的正长度或大小。半范数反而可以为非零的向量赋予零长度。

神经科学技术

神经科学,又称神经生物学,是专门研究神经系统的结构、功能、发育、演化、遗传学、生物化学、生理学、药理学及病理学的一门科学。对行为及学习的研究都是神经科学的分支。 对人脑研究是个跨领域的范畴,当中涉及分子层面、细胞层面、神经小组、大型神经系统,如视觉神经系统、脑干、脑皮层。

参数技术

在数学和统计学裡,参数(英语:parameter)是使用通用变量来建立函数和变量之间关系(当这种关系很难用方程来阐述时)的一个数量。

TensorFlow技术

TensorFlow是一个开源软件库,用于各种感知和语言理解任务的机器学习。目前被50个团队用于研究和生产许多Google商业产品,如语音识别、Gmail、Google 相册和搜索,其中许多产品曾使用过其前任软件DistBelief。

张量技术

张量是一个可用来表示在一些矢量、标量和其他张量之间的线性关系的多线性函数,这些线性关系的基本例子有内积、外积、线性映射以及笛卡儿积。其坐标在 维空间内,有 个分量的一种量,其中每个分量都是坐标的函数,而在坐标变换时,这些分量也依照某些规则作线性变换。称为该张量的秩或阶(与矩阵的秩和阶均无关系)。 在数学里,张量是一种几何实体,或者说广义上的“数量”。张量概念包括标量、矢量和线性算子。张量可以用坐标系统来表达,记作标量的数组,但它是定义为“不依赖于参照系的选择的”。张量在物理和工程学中很重要。例如在扩散张量成像中,表达器官对于水的在各个方向的微分透性的张量可以用来产生大脑的扫描图。工程上最重要的例子可能就是应力张量和应变张量了,它们都是二阶张量,对于一般线性材料他们之间的关系由一个四阶弹性张量来决定。

测试驱动开发技术

测试驱动开发,英文全称Test-Driven Development,简称TDD,是一种不同于传统软件开发流程的新型的开发方法。 它要求在编写某个功能的代码之前先编写测试代码,然后只编写使测试通过的功能代码,通过测试来推动整个开发的进行。 这有助于编写简洁可用和高质量的代码,并加速开发过程。

MXNet技术

MXNet是开源的,用来训练部署深层神经网络的深度学习框架。它是可扩展的,允许快速模型训练,并灵活支持多种语言(C ++,Python,Julia,Matlab,JavaScript, Go,R,Scala,Perl,Wolfram语言)

推荐文章
暂无评论
暂无评论~