Взаимодействие с NumPy

Протокол интерфейса массива

Протокол интерфейса массива определяет способ, которым объекты, похожие на массивы, могут повторно использовать буферы данных друг друга. Его реализация зависит от наличия следующих атрибутов или методов:

• __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 протокол определяет структуру памяти смещённых многомерных объектов массивов. Он предлагает следующий синтаксис для обмена данными:

1. Функция numpy.from_dlpack, которая принимает объекты (массивы) с методом __dlpack__ и использует этот метод для создания нового массива, содержащего данные из x.
2. Методы __dlpack__(self, stream=None) и __dlpack_device__ объекта массива, которые будут вызываться изнутри from_dlpack, для запроса устройства, на котором находится массив (может потребоваться передать правильный поток, например, в случае нескольких GPU) и для доступа к данным.

В отличие от протокола буфера, DLPack позволяет обмениваться массивами, содержащими данные на устройствах, отличных от процессора (например, Vulkan или GPU). Поскольку NumPy поддерживает только процессор, он может преобразовывать только объекты, данные которых находятся на процессоре. Но другие библиотеки, такие как PyTorch и CuPy, могут обмениваться данными на GPU с помощью этого протокола.

Протокол __array_ufunc__

Универсальная функция (или ufunc) — это «векторизованная» оболочка для функции, которая принимает фиксированное количество определённых входных данных и производит фиксированное количество определённых выходных данных. Выход ufunc (и её методов) не обязательно является многомерным массивом, если не все входные аргументы являются многомерными массивами. Действительно, если какой-либо входной параметр определяет метод __array_ufunc__, управление будет полностью передано этой функции, то есть ufunc переопределяется. Метод __array_ufunc__, определённый на этом (не многомерном массиве) объекте, имеет доступ к NumPy ufunc. Поскольку ufunc имеют чётко определённую структуру, внешний метод __array_ufunc__ может полагаться на атрибуты ufunc, такие как .at(), .reduce() и другие.

Подкласс может переопределить действия при выполнении NumPy ufuncs на нём, переопределив метод по умолчанию ndarray.__array_ufunc__. Этот метод выполняется вместо ufunc и должен возвращать либо результат операции, либо NotImplemented, если запрошенная операция не реализована.

Протокол __array_function__

Для обеспечения достаточного охвата API NumPy для проектов нижнего уровня необходимо перейти за пределы __array_ufunc__ и реализовать протокол, позволяющий аргументам функции NumPy взять под контроль и перенаправить выполнение в другую функцию (например, в GPU или параллельную реализацию) безопасным и согласованным способом для разных проектов.

Семантика __array_function__ очень похожа на __array_ufunc__, за исключением того, что операция указывается произвольным вызываемым объектом, а не экземпляром ufunc и методом. Более подробную информацию см. в NEP 18 — Механизм диспетчеризации для функций NumPy высокого уровня для массивов.

Пример: объекты Pandas

Рассмотрим следующее:

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

Теперь, ser это не ndarray, но поскольку он реализует протокол __array_ufunc__, мы можем применять ufunc к нему, как если бы это был ndarray:

>>> 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

Мы можем даже выполнять операции с другими ndarray:

>>> 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 обычно называются тензорами. Тензоры похожи на ndarray NumPy, за исключением того, что тензоры могут работать на 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)

Также обратите внимание, что возвращаемый тип этой функции совместим с исходным типом данных.

Обратите также внимание, что мы можем увидеть __array_wrap__ в действии здесь, даже если torch.Tensor не является подклассом ndarray:

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

PyTorch реализует __array_wrap__ для возможности получения тензоров из функций NumPy, и мы можем напрямую изменить его, чтобы контролировать, какие типы объектов возвращаются из этих функций.

Пример: массивы CuPy

CuPy — это совместимая с NumPy/SciPy библиотека массивов для ускоренных вычислений на GPU с Python. CuPy реализует подмножество интерфейса NumPy, реализуя cupy.ndarray, аналог NumPy ndarray.

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

Объект cupy.ndarray реализует интерфейс __array_ufunc__. Это позволяет применять NumPy ufunc к массивам CuPy (это делегирует операцию соответствующей CuPy CUDA/ROCm реализации ufunc):

>>> 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 реализует многие функции NumPy для объектов cupy.ndarray, но не все. Подробности см. в документации 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 и объем взаимодействия массивов 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

>>> 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.

>>> 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