我希望有人告诉我有关Tensor计算库的信息

2020-12-16 05:02:50

在此博客文章中,我们将分解张量计算库的实际含义以及它们之间的区别。我们将详细介绍一些流行的库,最后以当代张量计算库为背景对Theano的未来进行观察。

我对张量计算库(或计算图库,或符号代数库,或当今它们正在推销的任何内容)感到困惑。

我最初被介绍给PyTorch和TensorFlow,并且在没有其他参考的情况下认为它们是张量计算库的典型示例。然后我了解了Theano-一个较旧且流行的项目,但与PyTorch和TensorFlow有所不同,并且在一些有意义的方式上有所改进.JAX之后是它,它基本上是NumPy,带有更多的花哨功能(尽管我无法确切说明他们是)。然后,PyMC开发人员宣布Theano将拥有一个新的JAXbackend。

与我之前关于概率编程框架剖析的帖子类似,我将首先讨论一般的张量计算库-它们是什么以及它们之间的区别。然后,我将详细讨论一些库,最后在当代张量计算库的背景下对Theano的未来进行观察。

他们运行计算本身(duh),但也运行“相关”计算,它们要么(a)使用计算图,要么(b)直接对计算图本身进行操作,后者的一个很好的例子是优化计算本身:思考符号简化(例如xy / x = y)或数字稳定性的修改(例如,对于小x值,为log(1 + x))。

而且它们为计算提供了“最佳执行”:无论是通过JIT(即时)编译来更改执行,通过利用特殊硬件(GPU / TPU),通过对计算进行矢量化还是以任何其他方式来改变执行。

顺便说一句:我意识到“张量计算库”的名称太宽泛,并且上面的特征排除了一些可能被合理地称为“张量计算库”的库。更好的名称可能是“图形计算库”(尽管可能与networkx之类的库混在一起)或“计算图管理库”,甚至是“符号张量代数库”。

因此,为避免疑问,以下是此博客文章所不涉及的库列表:

NumPy和SciPy这些库没有计算图的概念-它们更像是一个函数工具箱,从Python调用并在C或Fortran中执行。

但是,这可能是一个有争议的区别-正如我们稍后将要看到的,JAX也没有建立显式的计算图,我当然想将JAX包含为“张量计算库”…\ _(ツ)_ /

Numba和Cython这些库为代码提供了最佳执行(实际上,一些张量计算库(例如Theano)很好地利用了它们),但是像NumPy和SciPy一样,它们实际上并不管理计算图本身。

Keras,Trax,Flax和PyTorch-Lightning这些库是张量计算库的高级包装-它们基本上提供抽象和面向用户的API,以更友好的方式利用张量计算库。

上述所有三个目标都是具有复杂解决方案的雄心勃勃的事业,因此,了解追求目标的决定可能会对其他目标产生影响(甚至要进行权衡!)也就不足为奇了。以下是所有三个轴上的共同差异列表:

张量计算库在如何表示计算图以及如何构建方面可以有所不同。静态图或动态图:首先要完全定义图,然后将数据注入以运行(又称为“定义并运行”),还是通过实际的正向计算即时定义图(即按运行来定义)? ? TensorFlow 1.x以其静态图而闻名,这使用户感到“正在通过锁孔处理计算图”,尤其是与PyTorch的动态图相比时。

懒惰或渴望执行:是在定义变量后立即评估变量,还是仅在评估自变量时才评估变量?通常,张量计算库要么选择支持急切执行的动态图,要么选择支持懒惰执行的静态图-例如,TensorFlow 2.0支持这两种模式。

有趣的是,某些张量计算库(例如Thinc)甚至没有构造明确的计算图:它们将其表示为链接的高阶函数。

Tensor计算库在使用计算图的目的上也可能有所不同-例如,我们的目标是基本上完成以“不同模式”运行计算图的事情,还是目标是自己修改计算图?显然,如何表示计算图以及要将其用于什么是非常相关的问题!例如,如果您希望能够将aribtrary计算形式表示为图形,则必须处理诸如if-else语句或for-loops之类的控制流-这导致在JAX中使用Python for-loops或需要使用手电筒导致常见问题。使用PyTorch在for循环中使用nn.ModuleList。

一些张量计算库(例如Theano和itfork的Theano-PyMC)旨在优化计算图形本身,为此必须有一个明确的图形。

最后,张量计算库在执行代码方面也可能有所不同。所有张量计算库都在CPU上运行,但是GPU和TPU支持的强度是张量计算库之间的主要区别。

另一个区别是张量计算库如何编译要在硬件上执行的代码。例如,他们是否使用JIT编译?他们是否使用“原始” C或CUDA编译器,或XLA编译器来进行机器学习的特定代码?

概述了张量计算库的基本相似之处和不同之处之后,我认为以几种流行的库为例会很有帮助。我试图尽可能地链接到相关文档。 1个

计算图如何表示和构建? PyTorch动态构建(并热切评估)一个明确的计算图。有关如何完成此操作的更多详细信息,请查看有关autogradmechanics的PyTorch文档。

有关PyTorch计算图的更多信息,请参见jdhao在PyTorch中有关计算图的介绍性博客文章。

计算图是用来做什么的?引用PyTorch文档,“ PyTorch是一个优化的张量库,用于使用GPU和CPU进行深度学习” –因此,主要重点是自动区分。

计算图如何表示和构建?无需构建显式的计算图来计算梯度,JAX只需提供grad()即可返回任何提供的函数的梯度函数。因此,从技术上讲,没有计算图的概念,只有纯(即无状态和无副作用)函数及其梯度。

当您计算前向通过时,PyTorch会建立一个图,然后在某个“结果”节点上调用向后(),然后使用结果节点相对于该中间节点的梯度来扩大图中的每个中间节点。另一方面,JAX使您可以将您的计算表达为Python函数,并通过用grad()对其进行转换,可以像计算函数一样为您提供一个可以评估的梯度函数-但是,与输出相比,它为您提供了相对于输出的梯度(默认)函数输入的第一个参数。

计算图是用来做什么的?根据JAX快速入门,JAX自称为“ NumPy在CPU,GPU和TPU上具有出色的自动差异性,可用于高性能机器学习研究”。因此,它的重点是自动分化。

该库如何确保计算的“最佳执行”? JAX使用XLA在[…] GPU和TPU上编译和运行您的NumPy代码。默认情况下,编译是在后台进行的,并且库调用会及时进行编译和执行。但是,JAX甚至可以让您及时将自己的Python函数编译为XLA优化的内核[…]编译和自动区分可以任意组合[…]

有关JAX的四功能API(grad,jit,vmap和pmap)的更多详细信息,请参阅Alex Minaar关于JAX的工作原理的概述。

注意:原始的Theano(由MILA维护)已经停产,并且PyMC开发人员已分叉了该项目:Theano-PyMC(即将重命名为Aesara)。我将在下面讨论原始项目和分支项目。

计算图是用来做什么的? Theano在张量计算库中是唯一的,因为它更加强调计算图本身的合理性。换句话说,尽管Theano对自动微分功能有强大的支持,但运行计算和计算梯度并不是一劳永逸的事情:Theano具有用于优化计算图形本身的anentire模块,并且可以很容易地将Theano图编译为不同的计算结果后端(默认情况下,Theano编译为C或CUDA,但是直接编译为JAX很简单)。

Theano通常被认为是进行深度学习研究的图书馆,但不仅如此!

该库如何确保计算的“最佳执行”?最初的Theano使用GCC C编译器进行CPU计算,并使用NVCC CUDA编译器进行GPU计算。

Theano-PyMC fork项目将使用JAX作为后端,可以利用可用的CPU,GPU和TPU。

最后,快速观察静态图和Theano填补的其他张量计算库所没有的生态位。在这一部分,我得到了Thomas Wiecki和Brandon Willard的大力帮助。

大多数张量计算库中的静态图形(或更确切地说,是静态生成的图形)始终保持一致:PyTorch和TensorFlow 2默认都支持动态生成的图形,而JAX完全放弃了显式的计算图形。

这种运动是可以理解的-动态地建立计算图可以更好地匹配人们的编程直觉。当我写z = x + y时,我的意思不是“我想用两个输入注册一个求和运算,正在等待数据注入”-我的意思是“我想计算x和y的和”。额外的间接层对大多数只希望以合理速度运行张量计算的用户没有帮助。

明确表示计算图对于某些事物非常有用,即使它使该图更难使用。您可以修改图形(例如图形优化,简化和重写),并且可以推理和分析图形。将计算作为实际对象,对您需要考虑计算本身而不是盲目运行的任务有不可估量的帮助。

另一方面,对于动态生成的图,计算图实际上从未在任何地方定义:计算是在运行中和在幕后进行的。您无法再对计算图做任何有趣的事情:例如,如果计算很慢,就无法推断出图的哪些部分很慢。最终结果是,您基本上必须希望框架内部执行正确的操作,而他们可能做不到!

这是Theano(或更确切地说,Theano-PyMC / Aesara)填补了其他当代张量计算库所不能满足的利基市场:承诺是,如果您花时间在前面指定您的计算并且全部一次,Theano可以优化生活日光无需进行任何计算-无论是通过图形操作,高效编译还是完全通过其他方式进行-只需执行一次即可。

一些读者会注意到,此列表中明显缺少TensorFlow-排除TensorFlow并非出于恶意,只是缺乏时间和精力来进行必要的研究以使其公正。抱歉。 ↩

订阅我的新闻通讯!每两周发表一次的关于阅读和学习内容的想法。这里有更多信息。这里有新闻通讯存档。