与 NumPy 的互操作性#

NumPy 的 ndarray 对象为数组结构化数据的操作提供了一个高级 API,并基于跨步内存存储提供了一个具体的 API 实现.虽然这个 API 功能强大且相当通用,但其具体实现有其局限性.随着数据集的增长和 NumPy 在各种新环境和架构中的使用,跨步内存存储策略在某些情况下是不合适的,这导致不同的库为了自己的用途重新实现了这个 API.这包括 GPU 数组 (CuPy)、稀疏数组 (scipy.sparse, PyData/Sparse) 和平行数组 (Dask arrays) 以及各种深度学习框架中的类似 NumPy 的实现,如 TensorFlow 和 PyTorch.同样,有许多项目在 NumPy API 的基础上构建,用于标记和索引数组 (XArray)、自动微分 (JAX)、掩码数组 (numpy.ma)、物理单位 (astropy.units, pint, unyt) 等,这些项目在 NumPy API 的基础上增加了额外的功能.

然而,用户仍然希望使用熟悉的 NumPy API 来处理这些数组,并在尽量减少（理想情况下为零）移植开销的情况下重用现有代码.为此,定义了各种协议,用于实现与 NumPy 匹配的高级 API 的多维数组.

广义上讲,有三个功能组用于与 NumPy 的互操作性:

将外部对象转换为 ndarray 的方法;
从 NumPy 函数延迟执行到另一个数组库的方法;
使用 NumPy 函数并返回一个外部对象实例的方法.

我们在下面描述这些特性.

在 NumPy 中使用任意对象#

NumPy API 的第一组互操作性特性允许在可能的情况下将外部对象视为 NumPy 数组.当 NumPy 函数遇到外部对象时,它们将按顺序尝试:

缓冲区协议,在 Python C-API 文档中有描述.
__array_interface__ 协议,在本页中描述.这是 Python 缓冲协议的前身,它定义了一种从其他 C 扩展访问 NumPy 数组内容的方法.
__array__() 方法,它要求一个任意对象将自己转换为一个数组.

对于缓冲区和 __array_interface__ 协议,对象描述其内存布局,NumPy 完成其余所有工作（如果可能,零拷贝）.如果不可能,对象本身负责从 __array__() 返回一个 ndarray.

DLPack 是另一种以语言和设备无关的方式将外部对象转换为 NumPy 数组的协议.NumPy 不会隐式地使用 DLPack 将对象转换为 ndarray.它提供了 numpy.from_dlpack 函数,该函数接受任何实现 __dlpack__ 方法的对象并输出一个 NumPy ndarray（这通常是输入对象数据缓冲区的视图）.:ref:dlpack:python-spec 页面详细解释了 __dlpack__ 协议.

数组接口协议#

数组接口协议定义了一种使类似数组的对象能够重用彼此数据缓冲区的方法.其实现依赖于以下属性或方法的存在:

__array_interface__: 一个包含数组形状、元素类型,以及可选的数组数据缓冲区地址和步幅的Python字典;
__array__(): 返回一个数组类对象的 NumPy ndarray 副本或视图的方法;

__array_interface__ 属性可以直接检查:

>>> import numpy as np
>>> x = np.array([1, 2, 5.0, 8])
>>> x.__array_interface__
{'data': (94708397920832, False), 'strides': None, 'descr': [('', '<f8')], 'typestr': '<f8', 'shape': (4,), 'version': 3}

__array_interface__ 属性也可以用来就地操作对象数据:

>>> class wrapper():
...     pass
...
>>> arr = np.array([1, 2, 3, 4])
>>> buf = arr.__array_interface__
>>> buf
{'data': (140497590272032, False), 'strides': None, 'descr': [('', '<i8')], 'typestr': '<i8', 'shape': (4,), 'version': 3}
>>> buf['shape'] = (2, 2)
>>> w = wrapper()
>>> w.__array_interface__ = buf
>>> new_arr = np.array(w, copy=False)
>>> new_arr
array([[1, 2],
       [3, 4]])

我们可以检查 arr 和 new_arr 共享相同的数据缓冲区:

>>> new_arr[0, 0] = 1000
>>> new_arr
array([[1000,    2],
       [   3,    4]])
>>> arr
array([1000, 2, 3, 4])

`array()` 方法#

__array__() 方法确保任何实现它的类 NumPy 对象（一个数组,任何暴露数组接口的对象,一个 __array__() 方法返回数组的对象或任何嵌套序列）都可以用作 NumPy 数组.如果可能,这将意味着使用 __array__() 创建数组类对象的 NumPy ndarray 视图.否则,这将数据复制到一个新的 ndarray 对象中.这不是最优的,因为将数组强制转换为 ndarray 可能会导致性能问题或需要复制和丢失元数据,因为原始对象及其可能具有的任何属性/行为都会丢失.

该方法的签名应为 __array__(self, dtype=None, copy=None).如果传入的 dtype 不是 None 并且与对象的数据类型不同,则应发生指定类型的转换.如果 copy 是 None,则只有在 dtype 参数强制要求时才应进行复制.对于 copy=True,应始终进行复制,而 copy=False 应在需要复制时引发异常.

如果一个类实现了旧的签名 __array__(self),对于 np.array(a) 将会引发一个警告,指出 dtype 和 copy 参数缺失.

要查看自定义数组实现的示例,包括使用 __array__(),请参见编写自定义数组容器.

DLPack 协议#

DLPack 协议定义了跨步 n 维数组对象的内存布局.它提供了以下数据交换语法:

numpy.from_dlpack 函数,接受带有 __dlpack__ 方法的（数组）对象,并使用该方法构造一个包含 x 数据的新数组.
__dlpack__(self, stream=None) 和 __dlpack_device__ 方法在数组对象上,这些方法将在 from_dlpack 内部调用,以查询数组所在的设备（可能需要传入正确的流,例如在多GPU的情况下）并访问数据.

与缓冲区协议不同,DLPack 允许交换包含在 CPU 之外设备上的数据的数组（例如 Vulkan 或 GPU）.由于 NumPy 仅支持 CPU,它只能转换数据存在于 CPU 上的对象.但其他库,如 PyTorch 和 CuPy,可能使用此协议在 GPU 上交换数据.

2. 在不转换的情况下操作外部对象#

NumPy API 定义的第二组方法允许我们将执行从 NumPy 函数延迟到另一个数组库.

考虑以下函数.

>>> import numpy as np
>>> def f(x):
...     return np.mean(np.exp(x))

注意 np.exp 是一个 ufunc ,这意味着它以逐元素的方式对 ndarrays 进行操作.另一方面,`np.mean <numpy.mean>` 沿着数组的一个轴进行操作.

我们可以直接对 NumPy ndarray 对象应用 f:

>>> x = np.array([1, 2, 3, 4])
>>> f(x)
21.1977562209304

我们希望这个函数能够同样好地处理任何类似 NumPy 的数组对象.

NumPy 允许一个类通过以下接口表明它希望以自定义定义的方式处理计算:

__array_ufunc__: 允许第三方对象支持和覆盖 ufuncs.
__array_function__:一个用于捕获未被通用函数 __array_ufunc__ 协议涵盖的 NumPy 功能的总称.

只要外部对象实现了 __array_ufunc__ 或 __array_function__ 协议,就可以在不进行显式转换的情况下对它们进行操作.

`__array_ufunc__` 协议#

一个通用函数（或简称 ufunc）是一个为函数提供的”矢量化”包装器,该函数接受固定数量的特定输入并产生固定数量的特定输出.ufunc 的输出（及其方法）不一定是 ndarray,如果并非所有输入参数都是 ndarray 的话.实际上,如果任何输入定义了 __array_ufunc__ 方法,控制权将完全传递给该函数,即,ufunc 被覆盖.在该（非 ndarray）对象上定义的 __array_ufunc__ 方法可以访问 NumPy ufunc.由于 ufuncs 具有定义良好的结构,外部的 __array_ufunc__ 方法可以依赖于 ufunc 属性,如 .at()、.reduce() 等.

子类可以通过重写默认的 ndarray.__array_ufunc__ 方法来覆盖在它上面执行 NumPy ufuncs 时的行为.该方法将代替 ufunc 执行,并应返回操作的结果,或者在请求的操作未实现时返回 NotImplemented.

`__array_function__` 协议#

为了实现对NumPy API的足够覆盖以支持下游项目,需要超越 __array_ufunc__ 并实现一个协议,该协议允许NumPy函数的参数接管并引导执行到另一个函数（例如,GPU或并行实现）,以一种在项目间安全且一致的方式.

__array_function__ 的语义与 __array_ufunc__ 非常相似,不同之处在于操作由任意可调用对象指定,而不是 ufunc 实例和方法.更多详情请参见 NEP 18 — A dispatch mechanism for NumPy’s high level array functions.

3. 返回外部对象#

第三种类型的功能集旨在使用 NumPy 函数实现,然后将返回值转换回外部对象的实例.``__array_finalize__`` 和 __array_wrap__ 方法在幕后工作,以确保 NumPy 函数的返回类型可以根据需要指定.

__array_finalize__ 方法是 NumPy 提供的一个机制,允许子类处理以各种方式创建新实例的情况.每当系统从 ndarray 的子类（子类型）对象内部分配新数组时,都会调用此方法.它可用于在构造后更改属性,或从”父类”更新元信息.

__array_wrap__ 方法”结束动作”的意思是允许任何对象（如用户定义的函数）设置其返回值的类型并更新属性和元数据.这可以看作是 __array__ 方法的反面.在实现 __array_wrap__ 的每个对象的末尾,此方法会在具有最高 array priority 的输入对象上调用,或者如果指定了输出对象,则在输出对象上调用.``__array_priority__`` 属性用于确定在返回对象的 Python 类型有多种可能的情况下返回哪种类型的对象.例如,子类可以选择使用此方法将输出数组转换为子类的实例并在返回数组给用户之前更新元数据.

有关这些方法的更多信息,请参见子类化 ndarray 和 ndarray 子类型的特定特性.

互操作性示例#

示例:Pandas `Series` 对象#

考虑以下:

>>> import pandas as pd
>>> ser = pd.Series([1, 2, 3, 4])
>>> type(ser)
pandas.core.series.Series

现在,``ser`` 不是一个 ndarray,但由于它实现了 __array_ufunc__ 协议 ,我们可以像对待 ndarray 一样对其应用 ufuncs:

>>> np.exp(ser)
   0     2.718282
   1     7.389056
   2    20.085537
   3    54.598150
   dtype: float64
>>> np.sin(ser)
   0    0.841471
   1    0.909297
   2    0.141120
   3   -0.756802
   dtype: float64

我们甚至可以对其他 ndarrays 进行操作:

>>> np.add(ser, np.array([5, 6, 7, 8]))
   0     6
   1     8
   2    10
   3    12
   dtype: int64
>>> f(ser)
21.1977562209304
>>> result = ser.__array__()
>>> type(result)
numpy.ndarray

示例:PyTorch 张量#

PyTorch 是一个使用 GPU 和 CPU 进行深度学习的优化张量库.PyTorch 数组通常被称为张量.张量类似于 NumPy 的 ndarrays,不同之处在于张量可以在 GPU 或其他硬件加速器上运行.事实上,张量和 NumPy 数组通常可以共享相同的底层内存,从而无需复制数据.

>>> import torch
>>> data = [[1, 2],[3, 4]]
>>> x_np = np.array(data)
>>> x_tensor = torch.tensor(data)

注意 x_np 和 x_tensor 是不同类型的对象:

>>> x_np
array([[1, 2],
       [3, 4]])
>>> x_tensor
tensor([[1, 2],
        [3, 4]])

然而,我们可以将 PyTorch 张量视为 NumPy 数组,而无需显式转换:

>>> np.exp(x_tensor)
tensor([[ 2.7183,  7.3891],
        [20.0855, 54.5982]], dtype=torch.float64)

此外,请注意此函数的返回类型与初始数据类型兼容.

警告

虽然这种混合使用 ndarrays 和 tensors 可能很方便,但不推荐这样做.它不适用于非 CPU 的 tensors,并且在极端情况下会有意外行为.用户应更倾向于显式地将 ndarray 转换为 tensor.

备注

PyTorch 没有实现 __array_function__ 或 __array_ufunc__.在底层,``Tensor.__array__()`` 方法返回一个视图形式的 NumPy ndarray 的张量数据缓冲区.详情请参见这个问题和 __torch_function__ 实现.

还要注意,即使 torch.Tensor 不是 ndarray 的子类,我们在这里也可以看到 __array_wrap__ 的作用:

>>> import torch
>>> t = torch.arange(4)
>>> np.abs(t)
tensor([0, 1, 2, 3])

PyTorch 实现了 __array_wrap__ 以便能够从 NumPy 函数中获取张量,我们可以直接修改它来控制这些函数返回的对象类型.

示例:CuPy 数组#

CuPy 是一个兼容 NumPy/SciPy 的数组库,用于使用 Python 进行 GPU 加速计算.CuPy 通过实现 cupy.ndarray 来实现 NumPy 接口的一个子集,`这是 NumPy ndarrays 的对应物 <https://docs.cupy.dev/en/stable/reference/ndarray.html>`__.

>>> import cupy as cp
>>> x_gpu = cp.array([1, 2, 3, 4])

cupy.ndarray 对象实现了 __array_ufunc__ 接口.这使得 NumPy ufuncs 可以应用于 CuPy 数组（这将把操作推迟到与 ufunc 匹配的 CuPy CUDA/ROCm 实现）:

>>> np.mean(np.exp(x_gpu))
array(21.19775622)

注意这些操作的返回类型仍然与初始类型一致:

>>> arr = cp.random.randn(1, 2, 3, 4).astype(cp.float32)
>>> result = np.sum(arr)
>>> print(type(result))
<class 'cupy._core.core.ndarray'>

请参阅 CuPy 文档中的此页面以获取详细信息.

cupy.ndarray 也实现了 __array_function__ 接口,这意味着可以进行诸如

>>> a = np.random.randn(100, 100)
>>> a_gpu = cp.asarray(a)
>>> qr_gpu = np.linalg.qr(a_gpu)

CuPy 在 cupy.ndarray 对象上实现了许多 NumPy 函数,但不是所有的.详情请参阅 CuPy 文档.

示例:Dask 数组#

Dask 是一个灵活的 Python 并行计算库.Dask Array 使用分块算法实现了 NumPy ndarray 接口的一个子集,将大数组分割成许多小数组.这允许在多核上对大于内存的数组进行计算.

Dask 支持 __array__() 和 __array_ufunc__.

>>> import dask.array as da
>>> x = da.random.normal(1, 0.1, size=(20, 20), chunks=(10, 10))
>>> np.mean(np.exp(x))
dask.array<mean_agg-aggregate, shape=(), dtype=float64, chunksize=(), chunktype=numpy.ndarray>
>>> np.mean(np.exp(x)).compute()
5.090097550553843

备注

Dask 是惰性求值的,计算的结果直到你通过调用 compute() 请求时才会计算.

详情请参见 Dask 数组文档和 Dask 数组与 NumPy 数组的互操作性范围.

示例:DLPack#

几个 Python 数据科学库实现了 __dlpack__ 协议.其中包括 PyTorch 和 CuPy.实现此协议的库的完整列表可以在 DLPack 文档的这一页找到.

将 PyTorch CPU 张量转换为 NumPy 数组:

>>> import torch
>>> x_torch = torch.arange(5)
>>> x_torch
tensor([0, 1, 2, 3, 4])
>>> x_np = np.from_dlpack(x_torch)
>>> x_np
array([0, 1, 2, 3, 4])
>>> # note that x_np is a view of x_torch
>>> x_torch[1] = 100
>>> x_torch
tensor([  0, 100,   2,   3,   4])
>>> x_np
array([  0, 100,   2,   3,   4])

导入的数组是只读的,因此就地写入或操作将会失败:

>>> x.flags.writeable
False
>>> x_np[1] = 1
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ValueError: assignment destination is read-only

为了对导入的数组进行就地操作,必须创建一个副本,但这意味着会复制内存.对于非常大的数组,请不要这样做:

>>> x_np_copy = x_np.copy()
>>> x_np_copy.sort()  # works

备注

请注意,GPU 张量不能转换为 NumPy 数组,因为 NumPy 不支持 GPU 设备:

>>> x_torch = torch.arange(5, device='cuda')
>>> np.from_dlpack(x_torch)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: Unsupported device in DLTensor.

但是,如果两个库都支持数据缓冲区所在的设备,则可以使用 __dlpack__ 协议（例如 PyTorch 和 CuPy）:

>>> x_torch = torch.arange(5, device='cuda')
>>> x_cupy = cupy.from_dlpack(x_torch)

同样地,一个 NumPy 数组可以转换为一个 PyTorch 张量:

>>> x_np = np.arange(5)
>>> x_torch = torch.from_dlpack(x_np)

只读数组不能被导出:

>>> x_np = np.arange(5)
>>> x_np.flags.writeable = False
>>> torch.from_dlpack(x_np)  
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File ".../site-packages/torch/utils/dlpack.py", line 63, in from_dlpack
    dlpack = ext_tensor.__dlpack__()
TypeError: NumPy currently only supports dlpack for writeable arrays

进一步阅读#

数组.接口
编写自定义数组容器
特殊属性和方法 (关于 __array_ufunc__ 和 __array_function__ 协议的详细信息)
子类化 ndarray （关于 __array_wrap__ 和 __array_finalize__ 方法的详细信息）
ndarray 子类型的特定特性（关于 __array_finalize__、__array_wrap__ 和 __array_priority__ 实现的更多细节）
NumPy 路线图:互操作性
PyTorch 与 NumPy 桥接的文档