Ускоряем работу python с numba10.01.2024 14:45

Привет, Хабр!

Numba — это Just-In-Time компилятор, который превращает ваш код на питоне в машинный код на лету. Это не просто мелкая оптимизация, а серьёзно ускорение.

Если вы знакомы с интерпретируемыми языками, вы знаете, что они обычно медленнее компилируемых из-за необходимости анализировать и исполнять код на лету. Но что, если бы вы могли получить лучшее из обоих миров? JIT-компиляция позволяет интерпретируемому языку, каким является питон, динамически компилировать части кода в машинный код, значительно ускоряя исполнение.

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

Вы продолжаете писать на Python, а Numba заботится о скорости.

Быстренько установим!

Откройте любимый терминал и просто через пип сделайте инсталл:

pip install numba

Numba установлена. Но прежде чем мы перейдем к коду, убедимся, что у вас также установлен NumPy, поскольку Numba часто используется вместе с ним для максимальной производительности:

pip install numpy

Начнем с чего-то простого.

Сначала импортируем необходимые библиотеки:

import numpy as np
from numba import jit

Определим простую функцию, которую мы хотим ускорить с помощью Numba. Мы будем использовать декоратор @jit для ускорения (о декораторах чуть позже):

@jit(nopython=True)
def sum_array(arr):
    result = 0
    for i in arr:
        result += i
    return result

Функцияsum_array, росто суммирует элементы в массиве. Декоратор @jit(nopython=True) говорит Numba компилировать эту функцию в машинный код, обходясь без вмешательства интерпретатора Python.

Создадим большой массив и протестируем нашу функцию:

large_array = np.arange(1000000)  # Большой массив от 0 до 999999
print(sum_array(large_array))  # Вызываем нашу функцию

Запускаем скрипт:

python hello_numba.py

voilà! Мы можем заметить значительное ускорение по сравнению с обычной функцией, особенно на больших массивах.

@jit декоратор

Декоратор @jit — сокращение от Just-In-Time, что означает компиляцию «прямо на лету». С помощью @jit вы можете указать Numba компилировать определенную функцию, превращая её в молниеносный машинный код.

Когда вы помечаете функцию декоратором @jit, Numba анализирует вашу функцию и определяет, как её можно оптимизировать. Она смотрит на типы переменных и операции, которые вы выполняете, и на основе этого создает оптимизированный машинный код. Этот процесс происходит во время первого вызова функции, так что первый запуск может занять чуть больше времени, но все последующие вызовы будут быстрыми.

Numba имеет несколько параметров для тонкой настройки поведения @jit:

cache: Если установлено в True, Numba будет кэшировать скомпилированный код, что ускорит его выполнение при последующих запусках.

Кэширование полезно для функций, которые вызываются многократно с одинаковыми типами данных (весь код будет с замером time, можете сравнить с ванильным питоновским кодом):

@jit(nopython=True, cache=True)
def cached_sum_squares(arr):
    result = 0
    for i in arr:
        result += i ** 2
    return result

start_time = time.time()
cached_sum_squares(large_array)
print("Numba (cached) time:", time.time() - start_time)

nogil: При установке в True позволяет функции выполняться без GIL Python, что может быть полезно для многопоточных приложений:

from threading import Thread

@jit(nopython=True, nogil=True)
def nogil_sum(arr):
    return np.sum(arr)

def thread_function():
    nogil_sum(large_array)

# стартуем в двух потоках
thread1 = Thread(target=thread_function)
thread2 = Thread(target=thread_function)

start_time = time.time()
thread1.start()
thread2.start()
thread1.join()
thread2.join()
print("Numba (nogil) time:", time.time() - start_time)

parallel: Если установлено в True, Numba попытается автоматически распараллелить циклы в функции:

@jit(nopython=True, parallel=True)
def parallel_sum(arr):
    return np.sum(arr)

start_time = time.time()
parallel_sum(large_array)
print("Numba (parallel) time:", time.time() - start_time)

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

import numpy as np
from numba import jit, njit

# определяем функцию для вычисления квадрата числа
@njit(inline='always')  # Используем inline='always' для инлайнинга функции
def square(x):
    return x * x

# определяем основную функцию, которая использует функцию square
@njit
def sum_of_squares(arr):
    result = 0
    for i in arr:
        result += square(i)  # вызываем инлайненную функцию
    return result

# Тестовые данные
large_array = np.arange(1000)

# Вызываем функцию и печатаем результат
print(sum_of_squares(large_array))

Функция square помечена декоратором @njit(inline='always'), что указывает numba всегда инлайнить эту функцию в любую другую функцию, которая её вызывает. Когда мы вызываем sum_of_squares, numba компилирует эту функцию, автоматически инлайня функцию square прямо в цикл

Чем точнее вы определите типы данных в вашей функции, тем лучше Numba сможет её оптимизировать. Глобальные переменные могут замедлить оптимизацию, так как Numba не всегда может точно определить их типы и значения.

@vectorize

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

Декоратор @vectorize в намбе позволяет создавать векторизованные функции, которые работают как универсальные функции NumPy (ufuncs). Это означает, что вы можете писать функции на чистом питоне, а затем намба скомпилирует их в высокопроизводительные машинные инструкции, которые автоматически работают с массивами данных.

Допустим, у нас есть функция, которая вычисляет что-то важное (например, гипотенузу треугольника по двум катетам). Вот как мы можем ускорить её с помощью @vectorize:

import numpy as np
from numba import vectorize, float64
import time

# функция Python для вычисления гипотенузы
def pythagorean_theorem(a, b):
    return np.sqrt(a**2 + b**2)

# векторизованная функция с Numba
@vectorize([float64(float64, float64)])
def numba_pythagorean_theorem(a, b):
    return np.sqrt(a**2 + b**2)

# создаем большие массивы данных
a = np.array(np.random.sample(1000000), dtype=np.float64)
b = np.array(np.random.sample(1000000), dtype=np.float64)

# измеряем время выполнения для обычной функции
start_time = time.time()
pythagorean_theorem(a, b)
print("Обычное Python время:", time.time() - start_time)

# измеряем время выполнения для векторизованной функции
start_time = time.time()
numba_pythagorean_theorem(a, b)
print("Numba @vectorize время:", time.time() - start_time)

numba_pythagorean_theorem работает намного быстрее обычной pythagorean_theorem благодаря векторизации и компиляции Numba.

@vectorize не ограничивается только созданием ufuncs, он также позволяет вам указывать типы входных и выходных данных для доп. оптимизации. Создадим функцию, которая принимает два числа и возвращает их произведение. Определим сигнатуры для разных типов данных, таких как целые числа и числа с плавающей точкой:

import numpy as np
from numba import vectorize

# определяем функцию с несколькими сигнатурами для разных типов данных
@vectorize(['int32(int32, int32)', 'int64(int64, int64)', 'float32(float32, float32)', 'float64(float64, float64)'])
def multiply(a, b):
    return a * b

# сздаем массивы данных разных типов
int32_arr = np.arange(10, dtype=np.int32)
int64_arr = np.arange(10, dtype=np.int64)
float32_arr = np.arange(10, dtype=np.float32)
float64_arr = np.arange(10, dtype=np.float64)

@generated_jit

@generated_jit позволяет создавать функции, генерирующие специализированный код в зависимости от типов входных аргументов. В отличие от обычного @jit, который просто компилирует функцию для заданных типов данных, @generated_jit позволяет написать свою логику для выбора или генерации кода в зависимости от типов аргументов во время выполнения.

Когда вы вызываете функцию, декорированную @generated_jit, намба вызывает генератор функции — это ваш код, который определяет, какая версия функции должна быть скомпилирована и выполнена на основе типов аргументов. Это позволяет писать очень гибкие и оптимизированные функции.

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

from numba import generated_jit, types

@generated_jit
def smart_function(x):
    if isinstance(x, types.Integer):
        # Врсия функции для целых чисел
        def int_version(x):
            return x * 2
        return int_version
    elif isinstance(x, types.Float):
        # Версия функции для чисел с плавающей точкой
        def float_version(x):
            return x / 2
        return float_version

smart_function(10)   # Должно вернуть 20
smart_function(10.5)  # Должно вернуть 5.25

smart_function ведет себя по-разному в зависимости от типа входного аргумента: если это целое число, она удваивает его; если число с плавающей точкой — делит пополам.

@stencil

@stencil позволяет определять функции, которые автоматически применяются к каждому элементу массива с учетом его соседей. Оч.полезно для операций, которые зависят от локального контекста в массиве, например, для фильтрации или свертки.

С@stencil, вы определяете ядро— небольшую функцию, которая описывает вычисление на одном элементе и его соседях. Numba затем автоматически применяет это ядро ко всем элементам входного массива, создавая новый массив с результатами.

Допустим, мы хотим применить простой фильтр Собеля для выделения границ на изображении. Фильтр Собеля — это операция свертки, которая использует два ядра, одно для горизонтальных изменений, другое для вертикальных.

import numpy as np
from numba import stencil

# функциюя фильтра Собеля
@stencil
def sobel_kernel(a):
    return (a[-1, -1] - a[1, -1]) + 2 * (a[-1, 0] - a[1, 0]) + (a[-1, 1] - a[1, 1])

# сздаем тестовое изображение (массив)
image = np.random.rand(100, 100)

# фильтр Собеля
filtered_image = sobel_kernel(image)

# Результат — новый массив с примененным фильтром
print(filtered_image)

sobel_kernel определяет, как каждый пиксель изображения должен быть изменен на основе значений его соседей. @stencil автоматически обрабатывает границы и применяет это ядро ко всем пикселям входного изображения.

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

В завершение хочу порекомендовать бесплатные вебинары курса Python Developer. Professional, на которые могут зарегистрироваться все желающие:

© Habrahabr.ru