近年来,Python 的数组计算生态系统已有机地发展起来,以支持 GPU、稀疏数组和分布式数组。这很棒,也是去中心化开源开发中可能发生的增长的一个极佳例子。
然而,为了巩固这种增长并将其应用于整个生态系统,我们现在需要进行一些集中规划,以便从软件包需要相互了解的成对模型转变为软件包可以通过开发和遵守社区标准协议进行协调的生态系统模型。
通过适度的努力,我们可以定义 Numpy API 的一个子集,该子集在所有这些方面都能很好地工作,从而使生态系统能够在硬件之间更顺畅地过渡。本文描述了实现这一目标的机遇和挑战。
我们首先讨论两类库
Numpy 数组是数字 Python 生态系统的基础之一,并作为其他语言中类似库的标准模型。如今,它被用于分析卫星和生物医学图像、金融模型、基因组、海洋和大气、超级计算机模拟以及来自数千个其他领域的数据。
然而,Numpy 是在几年前设计的,其实现对于某些现代硬件来说不再是最优的,特别是多核工作站、多核 GPU 和分布式集群。
幸运的是,其他库在这些其他架构上实现了 Numpy 数组 API
因此,即使 Numpy 实现不再理想,**Numpy API** 仍然在后续项目中继续存在。
注意:Numpy 实现大多数时候仍然是理想的。密集内存数组仍然是常见情况。这篇博文讨论的是 Numpy 不理想的少数情况
因此今天我们可以编写在 Numpy、GPU、稀疏和并行数组之间类似的代码
import numpy as np
x = np.random.random(...) # 在单个 CPU 上运行
y = x.T.dot(np.log(x) + 1)
z = y - y.mean(axis=0)
print(z[:5])
import cupy as cp
x = cp.random.random(...) # 在 GPU 上运行
y = x.T.dot(cp.log(x) + 1)
z = y - y.mean(axis=0)
print(z[:5].get())
import dask.array as da
x = da.random.random(...) # 在多个 CPU 上运行
y = x.T.dot(da.log(x) + 1)
z = y - y.mean(axis=0)
print(z[:5].compute())
...
此外,每个深度学习框架(TensorFlow、PyTorch、MXNet)都有一个类似 Numpy 的东西,它与 Numpy 的 API **非常相似**,但绝对不是试图完全匹配。
在为不同硬件开发 Numpy API 的同时,许多库今天在 Numpy API 的基础上构建算法功能
这些项目等等增强了 Python 中的数组计算,在 Numpy 本身提供的功能之外构建了新功能。
还有一些项目,如 Pandas、Scikit-Learn 和 SciPy,它们使用 Numpy 的内存内部表示。在这篇博文中,我们将忽略这些库,而专注于那些只使用高级 Numpy API 而不使用低级表示的库。
考虑到这两类项目
我们希望将它们结合使用,例如将 Autograd 应用于 CuPy,将 TensorLy 应用于 Sparse,等等,包括所有未来可能出现的实现。这具有挑战性。
不幸的是,虽然所有数组实现的 API 都与 Numpy 的 API **非常相似**,但它们使用了不同的函数。
>>> numpy.sin is cupy.sin
False
这给使用方库带来了问题,因为它们现在需要根据接收到的类似数组的对象来切换使用的函数。
def f(x)
if isinstance(x, numpy.ndarray)
return np.sin(x)
elif isinstance(x, cupy.ndarray)
return cupy.sin(x)
elif ...
今天,每个数组项目都实现了一个自定义插件系统,用于在一些数组选项之间切换。如果你感兴趣,下面提供了这些插件机制的链接
例如,XArray 可以使用 Numpy 数组或 Dask 数组。这极大地惠及了该项目的用户,他们今天可以无缝地从笔记本电脑上的小型内存数据集过渡到集群上的 100TB 数据集,所有这些都使用相同的编程模型。然而,在考虑将稀疏数组或 GPU 数组添加到 XArray 的插件系统时,很快就清楚地看到,目前这将是昂贵的。
构建、维护和扩展这些插件机制是**昂贵的**。每个项目中的插件系统都不一样,因此任何新的数组实现都必须到每个库中多次构建相同的代码。同样,任何新的算法库都必须为每个 ndarray 实现构建插件。每个库都必须明确导入和理解其他库,并且必须随着这些库随时间变化而适应。这种覆盖并不完整,因此用户对其应用程序是否能够在硬件之间移植缺乏信心。
成对的插件机制对于单个项目来说是有意义的,但对于整个生态系统来说并不是一个高效的选择。
我今天看到两种解决方案
每种方法都有挑战。
我们可以构建一个新的库,这里称为 arrayish,它包含所有相关 Numpy 函数的可调度版本。然后我们说服所有人内部使用它而不是 Numpy。
因此,在每个类似数组的库的代码库中,我们编写如下代码
# 在 numpy 的代码库中
import arrayish
import numpy
@arrayish.sin.register(numpy.ndarray, numpy.sin)
@arrayish.cos.register(numpy.ndarray, numpy.cos)
@arrayish.dot.register(numpy.ndarray, numpy.ndarray, numpy.dot)
...
# 在 cupy 的代码库中
import arrayish
import cupy
@arrayish.sin.register(cupy.ndarray, cupy.sin)
@arrayish.cos.register(cupy.ndarray, cupy.cos)
@arrayish.dot.register(cupy.ndarray, cupy.ndarray, cupy.dot)
...
Dask、Sparse 以及任何其他类似 Numpy 的库等等,都依此类推。
在所有算法库(如 XArray、autograd、TensorLy 等)中,我们使用 arrayish 而不是 Numpy
# 在 XArray 的代码库中
# import numpy
import arrayish as numpy
这与之前的插件解决方案相同,但现在我们构建了一个社区标准的插件系统,希望所有项目都能同意使用。
这将维护多个插件系统的 n 乘 m 的巨大成本降低到在每个库中使用一个插件系统的更易于管理的 n 加 m 的成本。这个集中式项目也许还会受益于比任何单独项目更容易获得更好的维护。
然而,这也有成本
Hameer Abbasi 在此构建了一个 arrayish 的初步原型:github.com/hameerabbasi/arrayish。关于此话题,以 XArray+Sparse 为例,在pydata/sparse #1 中有一些讨论。
或者,中央调度机制可以位于 Numpy 本身内部。
Numpy 函数可以学习将控制权交给它们的参数,允许数组实现尽可能地接管。这将使现有的 Numpy 代码能够在外部开发的数组实现上工作。
这有先例。**__array_ufunc__** 协议允许任何定义了 __array_ufunc__ 方法的类接管任何 Numpy ufunc,例如 np.sin 或 np.exp。Numpy 的 reduction 函数(如 np.sum)已经会在其参数上查找 .sum 方法,如果可能则会委派给它们。
一些数组项目,如 Dask 和 Sparse,已经实现了 __array_ufunc__ 协议。此外,CuPy 还有一个开放的 PR。这里有一个示例,清晰地展示了 Numpy 函数在 Dask 数组上的使用。
>>> import numpy as np
>>> import dask.array as da
>>> x = da.ones(10, chunks=(5,)) # 一个 Dask 数组
>>> np.sum(np.exp(x)) # 将 Numpy 函数应用于 Dask 数组
dask.array<sum-aggregate, shape=(), dtype=float64, chunksize=()> # 得到一个 Dask 数组
我建议所有与 Numpy API 兼容的数组项目都实现 __array_ufunc__ 协议。
这适用于许多函数,但不是全部。其他操作如 tensordot、concatenate 和 stack 在算法代码中经常出现,但此处未涵盖。
这个解决方案避免了上述 arrayish 解决方案的社区挑战。大家都习惯于遵从 Numpy 的决定,并且只需要重写相对较少的代码。
这种方法面临的挑战在于,从历史上看,Numpy 的发展速度比生态系统的其他部分要慢。例如,上面提到的 __array_ufunc__ 协议经过了多年的讨论才被合并。幸运的是,Numpy 最近获得了资金,以帮助它更快地进行此类更改。不过,通过这项资助招聘的全职开发人员才刚刚开始工作,目前尚不清楚这项工作对他们来说有多重要。
无论如何,我更倾向于看到这种 Numpy 协议解决方案得以实现。
近年来,Python 的数组计算生态系统已有机地发展起来,以支持 GPU、稀疏数组和分布式数组。这很棒,也是去中心化开源开发中可能发生的增长的一个极佳例子。
然而,为了巩固这种增长并将其应用于整个生态系统,我们现在需要进行一些集中规划,以便从软件包需要相互了解的成对模型转变为软件包可以通过开发和遵守社区标准协议进行协调的生态系统模型。
社区之前已经经历过这种过渡(Numeric + Numarray -> Numpy,Scikit-Learn 的 fit/predict API 等),通常会取得令人惊讶的积极成果。
我今天面临的未决问题如下
在 BIDS 举行的五月 NumPy 开发者冲刺中讨论此话题后,我们中的一些人起草了一份 Numpy 改进提案 (NEP),可在此处获取。