Itamar Turner-Trauring

Автор Python-профайлера Sciagraph

Векторизация NumPy ускоряет операции с объёмными данными за счёт быстрого низкоуровневого кода. На основе NumPy, чтобы предоставить столь же быструю функциональность, написана и Pandas. Но векторизация — не панацея. Иногда проблема решается за счёт памяти, нужная операция не поддерживается или просто не нужна. У каждой из этих проблем есть разные решения.

Слово «векторизация» имеет несколько значений, я расскажу о низкоуровневых циклах

Если резюмировать подробное объяснение векторизации, в контексте Python у векторизации три значения. Это:

1. API для работы с массивными данными: например, код arr += 1 прибавит 1 к каждому элементу массива NumPy.

2. Низкоуровневый API, например, на C или Rust. Он быстро работает с большим объёмом данных — это основная тема статьи.

3. Инструкции SIMD, ускоряющие операции низкого уровня ещё сильнее. Подробности — в основном описании векторизации.

Для начала предположим, что работаем с Python по умолчанию, то есть CPython.

Пример: прибавление числа к целым числам по списку

from time import time

l = list(range(100_000_000))

start = time()
for i in range(len(l)):
    l[i] += 17
print("Elapsed: ", time() - start)

Список целых чисел Python в CPython — это серия указателей на довольно большие объекты. С точки зрения CPython это обычные объекты Python, и прибавление целого числа к каждому элементу списка состоит из следующих шагов:

1. Поиск типа каждого элемента в списке.

2. Поиск функции сложения данного типа чисел.

3. Вызов найденной функции с исходным и добавляемым объектами.

4. Конвертирование обоих объектов Python в машинные целые числа.

5. Сложение машинных чисел.

6. Оборачивание полученного целого в объект Python.

7. Поиск следующего элемента в списке Python, что само по себе — дорогая операция, включающая поиск методов итератора для типа диапазона, затем поиск операции индексации для типа списка, преобразование объекта индекса в машинное целое число… Много работы!

Напротив, массив целых чисел NumPy — это массив машинных целых чисел. Таким образом, добавление целого числа к каждому элементу — всего лишь несколько процессорных инструкций для каждого элемента; после инициализации это ничем не отличается от эквивалента на C.

Вот код NumPy:

from time import time
import numpy as np

l = np.array(range(100_000_000), dtype=np.uint64)

start = time()
l += 17
print("Elapsed: ", time() - start)

NumPy работает намного быстрее:

Ограничения векторизации

Векторизация ускоряет код, это здорово… но решение не идеальное. Вот первая проблема:

Большие бесполезные выделения памяти

Допустим, нужно узнать среднее расстояние элементов массива от числа 0. Один из способов сделать это — вычислить абсолютное значение каждого элемента и взять среднее значение элементов:

imoprt numpy as np

def mean_distance_from_zero(arr):
    return np.abs(arr).mean()

Это работает, и оба вычисления выполняются быстро, но нужен совершенно новый промежуточный массив абсолютных значений. Если исходный массив содержит 1000 МБ, появится лишняя 1000 МБ. Исходный массив можно перезаписать, чтобы сэкономить память, но он может понадобиться по другим причинам.

Чтобы вместо N копий осталась только одна дополнительная копия, можно воспользоваться операциями на месте. Есть библиотеки, например numexpr и Dask Array, способные с сокращением нагрузки на память выполнять пакетные или более интеллектуальные обновления на месте. А можно просто рассчитать всё поэлементно, собственным циклом, но это приведёт ко второй проблеме:

Ускоряется только то, что поддерживается

Чтобы векторизация работала, нужен низкоуровневый машинный код, который выполняет операцию и управляет циклом над данными. При переключении обратно, на циклы и функциональность обычного Python, скорость теряется. Например, очевидный способ реализовать mean_distance_from_zero() без увеличения потребляемой памяти выглядит так:

def mean_distance_from_zero(arr):
    total = 0
    for i in range(len(arr))
        total += abs(arr[i])
    return total / len(arr)

Но мы вернулись к циклу и операциям в Python, код снова медленный — и вот одна из причин, почему NumPy приходится так много реализовывать самой: код замедляется при каждом возврате к Python. В компилируемом языке есть вероятность, что компилятор найдёт способ оптимизировать вторую реализацию, и это даст ускорение, которое невозможно с серией отдельных векторизованных операций.

Векторизация ускоряет только массовые операции

Иногда ваш код работает медленно потому, что выполняет одну и ту же операцию со многими элементами данных одного и того же типа — тогда и пригодится векторизация. В других случаях расчёт слишком медленный по другой причине, а значит, векторизация не поможет.

Другие решения

Посмотрим три основных:

1. Ускоренный интерпретатор Python, PyPy — реализация CPython намного умнее: она может использовать JIT-компиляцию, ускоряя выполнение генерацией специализированного машинного кода.

2. Генерация машинного кода внутри CPython через Numba.

3. Компиляция кастомного кода с помощью Cython, C, C++ или Rust.

PyPy: ускоренный интерпретатор Python

С Python-списком целых чисел CPython на самом деле не делает ничего хитрого, поэтому и нужны массивы NumPy. PyPy здесь разумнее CPython: для такого случая он генерирует специализированный код.

Вернёмся к исходному примеру:

from time import time

l = list(range(100_000_000))

start = time()
for i in range(len(l)):
    l[i] += 17
print("Elapsed: ", time() - start)

И сравним Cython с PyPy:

$ python add_list.py
Elapsed:  6.125197410583496
$ pypy add_list.py
Elapsed:  0.11461925506591797

Скорость PyPy сравнима со скоростью NumPy — а это обычный цикл Python.

К сожалению, PyPy не особенно хорошо работает с NumPy; наивный невекторизованный цикл над массивом NumPy (например, реализованная выше mean()) в PyPy намного медленнее, чем в CPython. Поэтому PyPy бесполезен в попытке справиться с недостающим функционалом NumPy, а медленный невекторизованный цикл он даже замедлит ещё сильней. Но, если вы имеете дело с кодом на стандартных объектах Python, PyPy может оказаться намного быстрее CPython, поэтому лучше всего он подходит для ускорения, когда NumPy или Pandas не используются вообще.

Numba: JIT-функции, работающие с NumPy

Numba также выполняет JIT-компиляцию, но в отличие от PyPy действует как дополнение к CPython и создан для работы с NumPy. Посмотрим, как это решит наши проблемы:

Расширение NumPy через Numba

Недостающие операции для Numba не проблема: просто напишите собственную функцию. Вот, например, mean_distance_from_zero() на чистом Python и Numba:

from time import time
import numpy as np
from numba import njit

# Pure Python version:
def mean_distance_from_zero(arr):
    total = 0
    for i in range(len(arr))
        total += abs(arr[i])
    return total / len(arr)

# A fast, JITed version:
mean_distance_from_zero_numba = njit(
    mean_distance_from_zero
)

arr = np.array(range(10_000_000), dtype=np.float64)

start = time()
mean_distance_from_zero(arr)
print("Elapsed CPython: ", time() - start)

for i in range(3):
    start = time()
    mean_distance_from_zero_numba(arr)
    print("Elapsed Numba:   ", time() - start)

Первый вызов Numba небыстрый: библиотека должна скомпилировать кастомный машинный код. После этого вызовы ускоряются и CPython сильно уступает Numba:

$ python cpython_vs_numba.py
Elapsed CPython:  1.1473402976989746
Elapsed Numba:    0.1538538932800293
Elapsed Numba:    0.0057942867279052734
Elapsed Numba:    0.005782604217529297

Поскольку мы можем писать собственные дополнительные функции для NumPy с помощью Numba, отсутствующие операции — не проблема. Это также означает, что из-за ограничений доступных массивов NumPy не нужно создавать временные массивы, поэтому Numba может помочь решить первые две проблемы, с которыми мы столкнулись.

Numba не так полезна в невекторизованных операциях, поскольку они не входят в задачи проекта.

Скомпилированный код Cython (или Rust, C, C++)

Мы увидели два примера JIT-компиляции: целый JIT-интерпретатор и единичные JIT-функции. Третий способ ускориться — предварительная компиляция. При помощи Cython, Rust или C вы сможете:

1. Создавать быстрые операции для NumPy через встроенную в Cython поддержку NumPy, rust-numpy в Rust или просто C API Python. Сама NumPy в основном написана на C, а её расширения — и на других языках, таких как Fortran или C++.

2. Для случаев без векторизации, чтобы писать быстрые расширения, опять же можно использовать эти языки.

Но, чтобы скомпилировать эти расширения перед запуском Python, придётся добавить много настроек конфигурации упаковки или настроить сборку. Cython умеет компилировать при импорте, но это осложняет распространение вашего ПО: его пользователям будет нужен компилятор. Numba таких ограничений не имеет.

Выбор решения

С векторизованной операцией PyPy не поможет. Хотя сделать последний полезнее могут проекты типа HPy Project.

Остаются Numba и скомпилированные расширения.

• Компилировать нужно заранее, а это требует компилятора и сложной упаковки, поэтому обычно легче настроить Numba.

• Numba без дополнительной работы создаёт разные версии кода для разных типов: целых чисел или чисел с плавающей точкой; с компилируемым кодом вам нужно скомпилировать версию для каждого типа числа и, возможно, также для каждого типа написать код. С помощью Rust или C++ необходимость дублирования кода можно обойти дженериками и шаблонами.

• Numba — это сильно ограниченное подмножество Python, так что отладка, связанная с его ограничениями, может оказаться затруднительной, а компилируемые языки более гибкие; Cython поддерживает почти весь Python, Rust — сложный язык с отличными инструментами, и т. д.

PyPy часто ускоряет невекторизованные операции без изменений кода; можно просто попробовать свой код и посмотреть его скорость. Иначе остаётся оптимизация существующего кода, или выбор подходящего компилируемого языка.

А пока Python ускоряется, мы поможем прокачать ваши навыки или с самого начала освоить профессию, актуальную в любое время:

• Профессия Fullstack-разработчик на Python

• Профессия Data Scientist

Выбрать другую востребованную профессию.