Автор: Антон Броиловский Редактура: Александр Наздрюхин, Андрей Шадриков

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.