Знакомимся с Numba

Python — (ну очень) высокоуровневый язык программирования. При работе с ним не нужно задумываться о том, как управлять памятью. Не нужно заботиться о том, чтобы имя переменной вдруг не начало указывать на объект другого типа. Не нужно даже думать о race condition’ах. Можно даже начать исполнять код, не написав его целиком: строчка за строчкой, как мы это делаем в Jupyter.

Одновременно с удобствами, в предыдущем абзаце мы на самом деле перечислили и причины, из-за которых Python медленный: фрагментация памяти, динамическая типизация, GIL, отсутствие компилируемости — это его зоны роста. Кажется, сейчас так принято говорить, чтоб не обидеть 😉.

Python любят за его выразительность и простоту. sum(x ** 2 for x in arr if x % 2 == 0) . Сейчас мы всего одной строчкой выразили намерение посчитать сумму квадратов всех четных элементов списка.

А не любят его за медлительность и большое потребление памяти. Хорошая новость в том, что с этим можно бороться. Есть способы попроще: например, использование numpy при работе с большими матрицами ускорит нас благодаря векторизации и нивелированию фрагментации памяти. А есть способы посложнее: например, написать свои биндинги на C/C++.

Сегодня мы рассмотрим Numba как один из относительно простых способов ускорить python-код. Благодаря ней мне удавалось ускорить расчеты иногда в десятки, а иногда и в сотни раз.

Как использовать?

Вам достаточно просто установить библиотеку при помощи pip (pip install numba). Затем импортировать декоратор @njit или @jit (from numba import njit), обернуть вашу функцию и, если повезет, получить из коробки ускорение в несколько раз!

Что делает numba?

Всего есть два режима работы numba, с помощью которых можно сократить время исполнения вашего кода на питоне:

1. Nopython. В этом режиме numba сначала скомпилирует функцию, чтобы выполнять её без участия интерпретатора, а также несколько оптимизирует набор низкоуровневых инструкций, чтобы код работал еще быстрее. Такой способ позволит ускорить любой метод. Например, функция расчета среднего значения в большом массиве заработает в десятки раз быстрее.
2. Object mode. Единственная оптимизация numba в этом режиме — она найдет все циклы и попытается их сначала вынести в отдельную функцию, а затем скомпилировать. Остальной код выполнится с помощью интерпретатора питона. Здесь можно достичь ускорения, но значительно меньшего в сравнении с первым вариантом.

Numba всегда пытается скомпилировать любую функцию, обернутую декоратором @jit в режиме nopython, а если не получится — обратится за помощью к object mode. Для гарантии работы режима nopython используется либо декоратор @njit, либо аргумент nopython=True у декоратора @jit, тогда если функцию скомпилировать не удастся — numba выдаст ошибку.

Когда использовать?

Если в коде есть много математических расчетов и циклов, с огромной вероятностью numba — ваш спаситель.

Как выбрать режим работы?

Конечно, было бы здорово всегда использовать режим nopython — все работает хорошо и быстро.

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

Хорошие новости: в numba реализовано большинство методов, правда, у некоторых нет параметра axis. Плохие новости: методы из scipy и opencv использовать не получится 😢

В остальных случаях придется обращаться к object mode. Здесь мы получим ускорение лишь при наличии в коде большого количества циклов, которые numba сможет скомпилировать (условия для успешной компиляции здесь аналогичны режиму nopython). Иначе время выполнения либо не изменится, либо увеличится.

Умеет ли numba работать параллельно?

Да! Добавление флага parallel=True в режиме nopython позволит распараллелить код внутри функции. Чтобы распараллелить цикл, нужно использовать prange вместо питоновского range'a. Будьте аккуратнее с этим режимом: каждый элемент массива в цикле должен обрабатываться независимо, иначе функция начнет работать некорректно.

Интересный факт: в prange цикле можно использовать операторы +=, -=, *= и /=, что позволяет считать статистики гораздо быстрее.

Какие еще есть полезные флаги?

1. nogil — можно снять GIL (это механизм блокирующий возможность работать нескольким потокам одновременно);
2. fastmath — можно ускорить математические расчеты, немного пожертвовав точностью;
3. cache — можно ускорить перекомпиляцию функции.

А где же минусы?

Главный минус numba — ее практически невозможно нормально дебажить. При работе с numba впервые можно сойти с ума в попытках понять, почему она падает в сотый раз…

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

Как себе помочь?

За годы работы с этой библиотекой удалось собрать несколько заметок, которые могут помочь при работе с ней:

1. Поскольку numba компилирует функцию, необходима строгая типизация! Если вы скомпилировали функцию при помощи двух аргументов float32, а потом вызвали ту же функцию с аргументами типа float64 — функция будет компилироваться заново.
2. Использование range (for i in range(len(obj)) вместо итерирования по объекту (for item in obj) работает быстрее.
3. Цикл в 95% случаев быстрее векторной операции (касается как range, так и prange).
4. Два вложенных цикла обычно быстрее одного цикла и векторной операции.
5. Три вложенных цикла практически всегда медленнее, чем два вложенных цикла и одна векторная операция.
6. Также будьте аккуратны с использованием prange внутри prange — вы можете встретить неожиданные замедления в непонятные моменты времени.

Чтобы не быть голословными, мы сделали коллаб с примерами 😊