Ускорьте вычисления на Python с помощью Numba

Python — это интерпретируемый, скриптовый язык, и интерпретатор CPython заявил себя не самым быстрым. Ускорить вычисления можно с помощью NumPy и SciPy. Но что если вы хотите реализовать новый алгоритм без использования кода на низкоуровневом языке? Для вычислений, в основном использующих массивы, библиотека Numba может значительно повысить производительность вашего кода. В этой статье мы расскажем, почему использование NumPy недостаточно для ускорения, и рассмотрим, как библиотека Numba ускоряет код на Python.

Когда NumPy уже не помогает

Допустим, у нас есть очень большой массив, и нам нужно найти ее монотонно возрастающую версию, такую, что значения только увеличиваются (это не то же самое, что сортировка). Пример:

[1, 2, 1, 3, 3, 5, 4, 6] -> [1, 2, 2, 3, 3, 5, 5, 6]

Простейшая реализация на Python:

def monotonically_increasing(a):
    max_value = 0
    for i in range(len(a)):
        if a[i] > max_value:
            max_value = a[i]
        a[i] = max_value

Но здесь есть одна проблема. NumPy выполняет расчеты без вызовов Python. Но так как функция имеет цикл, то мы потеряем в производительности, так как потеряем все преимущества, которые предоставляет эта библиотека.

Машинное обучение на Python

Код курса
PYML
Ближайшая дата курса
7 ноября, 2022
Длительность обучения
24 ак.часов
Стоимость обучения
45 000 руб.

Для 10,000,000-го NumPy-массива данная функция выполнялась 2.5 секунды. Но можно ли сделать лучше?

Numba может увеличить скорость вычислений

Numba — это JIT (just-in-time) компилятор, который предназначен как раз для кода с использованием циклов. Это именно то, что нам нужно.

В наш код на Python нужно добавить всего лишь несколько строк:

from numba import njit

@njit
def monotonically_increasing(a):
    max_value = 0
    for i in range(len(a)):
        if a[i] > max_value:
            max_value = a[i]
        a[i] = max_value

Мы добились времени выполнения 0.19 секунд, что почти в 13 раз быстрее. Уже неплохо.

На самом деле в NumPy есть функция, которая делает то же самое, — numpy.maximum.accumulate. При ее использовании время выполнения ограничивается 0.03 секунды.

Что используется Время выполнения, мс
Цикл Python 2560
Цикл с Numba 190
np.maximum.accumulate 30

Введение в Numba

Конечно, если есть нужная функция в NumPy или SciPy, как например с accumulate, то проблема решена. А если ее нет? При этом хочется добиться быстрых вычислений. В этмом случае Numba может предложить:

  • выполнять код, как в стандартном интерпретаторе Python, так и в компилируемой версии;
  • легко и быстро использовать циклы в коде.

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

Еще Numba позволяет выполнять код на GPU [1]. Вы можете обратиться у документации за разными примерами.

Ограничения Numba

В первый раз, когда вызывается функция, декорированная njit, нужно сгенерировать необходимый машинный код. Это занимает некоторое время. Например, мы можем использовать команду %time в IPython для измерения времени на выполнения декорированной функции:

In [1]: from numba import njit

In [2]: @njit
   ...: def add(a, b): a + b

In [3]: %time add(1, 2)
CPU times: user 320 ms, sys: 117 ms, total: 437 ms
Wall time: 207 ms

In [4]: %time add(1, 2)
CPU times: user 17 µs, sys: 0 ns, total: 17 µs
Wall time: 24.3 µs

In [5]: %time add(1, 2)
CPU times: user 8 µs, sys: 2 µs, total: 10 µs
Wall time: 13.6 µs

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

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

In [8]: %time add(1.5, 2.5)
CPU times: user 40.3 ms, sys: 1.14 ms, total: 41.5 ms
Wall time: 41 ms

In [9]: %time add(1.5, 2.5)
CPU times: user 16 µs, sys: 3 µs, total: 19 µs
Wall time: 26 µs

Итак, если вам нужна высокая скорость выполнения над массивами данных и в NumPy или SciPy нет нужной функции, то попробуйте использовать в своем проекте Numba.

 

А о том, как выполнять быстрый код для решения задач Machine Learning вы узнаете на наших образовательных курсах в лицензированном учебном центре обучения и повышения квалификации руководителей и ИТ-специалистов (менеджеров, архитекторов, инженеров, администраторов, Data Scientist’ов и аналитиков Big Data) в Москве:

Источники
  1. Numba for CUDA GPUs

Добавить комментарий

Поиск по сайту