В предишната статия разгледахме как можем да използваме нишки (threads) за конкурентно изпълнение на задачи и подобряване на производителността на програмите си. Въпреки това, поради глобалното заключване на интерпретатора (GIL) в Python, използването на нишки не винаги води до значително ускорение за CPU-интензивни задачи.
В тази статия ще разгледаме два мощни подхода за ускоряване на numerically-intensive Python код: векторизация с NumPy и компилиране в машинен код с Numba. Ще покажем как тези техники могат да доведат до драстични подобрения в скоростта на изпълнение, понякога от порядъци, без да се налага да пишем експлицитен паралелен или нискоуровнев код.
Векторизация с NumPy
NumPy е фундаментална библиотека за научни изчисления в Python, предлагаща ефективни многомерни масиви и богат набор от функции за работа с тях. Една от най-мощните особености на NumPy е способността му за векторизиране на операции.
Векторизацията се отнася до използването на операции над цели масиви вместо скаларни операции върху отделни елементи. Вместо да използваме експлицитни цикли в Python, можем да делегираме работата на оптимизирани функции в NumPy, които работят върху целите масиви наведнъж. Това не само води до по-кратък и по-четим код, но и до значително по-бързо изпълнение.
Нека илюстрираме това с прост пример. Да предположим, че искаме да изчислим скаларното произведение (dot product) на два вектора:
import numpy as np
def python_dot(a, b):
result = 0
for i in range(len(a)):
result += a[i] * b[i]
return result
def numpy_dot(a, b):
return np.dot(a, b)
# Пример употреба
a = np.random.random(1000)
b = np.random.random(1000)
%timeit python_dot(a, b)
# 100 loops, best of 3: 7.48 ms per loop
%timeit numpy_dot(a, b)
# 100000 loops, best of 3: 11.6 µs per loop
Тук имаме две функции, които изчисляват същия резултат: python_dot, която използва експлицитен Python цикъл, и numpy_dot, която използва вградената np.dot функция. Когато ги тестваме върху случайни вектори с 1000 елемента, виждаме огромна разлика в производителността – NumPy версията е около 650 пъти по-бърза!
Защо векторизираният код е толкова по-бърз? Основната причина е, че NumPy функциите са имплементирани на ниско ниво в оптимизиран C код, който работи директно върху целите масиви, без режийните разходи на Python интерпретатора и итерирането елемент по елемент.
Освен np.dot, NumPy предлага множество други векторизирани операции, като елементни аритметични операции (+, -, *, /), математически функции (np.sin, np.exp, np.log), статистически функции (np.mean, np.std, np.sum), булева индексация и маскиране, и много други. Когато е възможно, винаги трябва да се стремим да използваме тези вградени функции и да мислим на ниво масив вместо елемент.
Компилиране на Python функции с Numba
Въпреки че векторизацията с NumPy е мощен инструмент, не всички алгоритми могат лесно да се изразят във векторизирана форма. В някои случаи се налага да прибегнем до експлицитни цикли и итерации. За щастие, има и друг начин да ускорим такъв код – компилиране до машинен код с Numba.
Numba е open-source JIT (just-in-time) компилатор за Python, разработен от Anaconda, който може да превежда Python функции в бърз машинен код в реално време. Numba се фокусира върху подмножество от Python и NumPy специално за numerically-oriented code и може да постигне значителни ускорения с малко или никакви промени в оригиналния код.
Да вземем предишния пример с dot product, но този път да опитаме да го ускорим с Numba:
import numpy as np
from numba import njit
@njit
def numba_dot(a, b):
result = 0
for i in range(len(a)):
result += a[i] * b[i]
return result
# Пример употреба
a = np.random.random(1000)
b = np.random.random(1000)
%timeit python_dot(a, b)
# 100 loops, best of 3: 7.48 ms per loop
%timeit numba_dot(a, b)
# 100000 loops, best of 3: 2.88 µs per loop
Тук дефинираме нова функция numba_dot, която е идентична на python_dot, но е декорирана с @njit (nopython mode JIT). Това указва на Numba да компилира функцията до машинен код и да я изпълни директно, без да преминава през Python интерпретатора.
Резултатите са впечатляващи – компилираната с Numba функция е около 2600 пъти по-бърза от чистата Python версия и дори 4 пъти по-бърза от векторизираната NumPy версия! Нещо повече, постигаме това без да променяме самия алгоритъм и запазвайки ясната итеративна структура на оригиналния код.
Numba поддържа голяма част от стандартния езикови функции на Python, както и голяма част от NumPy API-то, включително масиви, математически функции и случайно генериране. Той също така предлага удобни начини за паралелизиране на кода чрез автоматично паралелизиране на циклите и векторизирани математически функции.
Някои важни особености и ограничения на Numba, за които трябва да имаме предвид:
- Numba работи най-добре с numerically-oriented code, който оперира върху NumPy масиви и скаларни типове. Други типове данни като списъци, речници и потребителски класове не се поддържат добре и могат да попречат на процеса на компилиране.
- Функциите, които искаме да компилираме, трябва да са декорирани с
@jitили@njitи да следват определени ограничения – например, без рекурсия, без вложени функции, без нестандартни контролни потоци. - Компилирането с Numba може да отнеме значително време при първото извикване на функцията (заради процеса на JIT компилация). Следващите извиквания обаче ще бъдат много по-бързи, тъй като ще използват вече компилирания код.
- Отстраняването на грешки в компилиран с Numba код може да бъде по-трудно, тъй като стандартните Python инструменти за дебъгване не работят. Numba предлага отделен дебъгер, но той е по-ограничен.
Въпреки тези ограничения, Numba е изключително мощен инструмент за ускоряване на numerically-intensive Python код с минимални усилия. В много случаи той може дори да елиминира нуждата от пренаписване на критични секции на код на по-ниско ниво езици като C или Fortran.
Обобщение и ключови изводи
В тази статия разгледахме два ключови подхода за ускоряване на numerically-intensive Python код:
- Векторизация с NumPy: Използване на операции над цели масиви вместо явни цикли и итерации. Това позволява делегиране на работата към оптимизиран компилиран код и може да доведе до огромни ускорения с минимална промяна в кода.
- Компилиране до машинен код с Numba: Автоматично JIT компилиране на обикновени Python функции чрез декориране с
@jitили@njit. Това може да доведе до производителност, подобна на C, като същевременно запазва високото ниво и четимостта на Python.
В реалните приложения често се използва комбинация от двата подхода – използва се векторизация с NumPy когато е възможно, а частите, които не могат лесно да се изразят във векторизирана форма, се компилират с Numba.
Важно е да се отбележи, че макар тези техники да са изключително ефективни за numerically-intensive код, те не са универсално решение за всички проблеми с производителността в Python. За операции, доминирани от входно-изходна дейност, мрежови операции или достъп до база данни, други подходи като асинхронно програмиране или използване на специализирани библиотеки може да са по-подходящи.
Независимо от това, векторизацията и компилацията са задължителни техники в арсенала на всеки Python програмист, занимаващ се с numerically-heavy задачи. С правилните инструменти и подход Python може да постигне впечатляваща производителност на par с традиционните компилирани езици, без да се жертва неговата простота и изразителност.
Надявам се, че тази статия е дала ценни прозрения за ускоряване на numerically-intensive Python код. Не се колебайте да експериментирате с тези техники в собствените си проекти и да се гмурнете по-дълбоко в екосистемата от високопроизводителни Python инструменти като NumPy, Numba, Cython, Dask и други.
