В статье «Ускоряем анализ данных в 180 000 раз с помощью Rust» показано, как неоптимизированный код на Python, после переписывания и оптимизации на Rust, ускоряется в 180 000 раз. Автор отмечает: «есть множество способов сделать код на Python быстрее, но смысл этого поста не в том, чтобы сравнить высокооптимизированный Python с высокооптимизированным Rust. Смысл в том, чтобы сравнить "стандартный-Jupyter-notebook" Python с высокооптимизированным Rust».
Возникает вопрос: какого ускорения мы могли бы достичь, если бы остановились на Python?
Под катом разработчик Сидни Рэдклифф* проходит путь профилирования и итеративного ускорения кода на Python, чтобы выяснить это.
*Обращаем ваше внимание, что позиция автора может не всегда совпадать с мнением МойОфис.
Повторяем исходные бенчмарки
Как и в упомянутой выше статье, мы используем M1 Macbook, и по тем же бенчмаркам получаем сопоставимые показатели:
Среднее время итерации исходного неоптимизированного кода, измеренное за 1000 итераций, — 35 мс. В оригинальной статье — 36 мс.
После полной оптимизации код на Rust ускорен в 180,081 раз. В оригинальной статье сообщается о 182 450-кратном ускорении.
Изначальный код на Python
Вот неоптимизированный код на Python из ранее упомянутой статьи.
from itertools import combinations
import pandas as pd
from pandas import IndexSlice as islice
def k_corrset(data, K):
all_qs = data.question.unique()
q_to_score = data.set_index(['question', 'user'])
all_grand_totals = data.groupby('user').score.sum().rename('grand_total')
# Inner loop
corrs = []
for qs in combinations(all_qs, K):
qs_data = q_to_score.loc[islice[qs,:],:].swaplevel()
answered_all = qs_data.groupby(level=[0]).size() == K
answered_all = answered_all[answered_all].index
qs_totals = qs_data.loc[islice[answered_all,:]] \
.groupby(level=[0]).sum().rename(columns={'score': 'qs'})
r = qs_totals.join(all_grand_totals).corr().qs.grand_total
corrs.append({'qs': qs, 'r': r})
corrs = pd.DataFrame(corrs)
return corrs.sort_values('r', ascending=False).iloc[0].qs
data = pd.read_json('scores.json')
print(k_corrset(data, K=5))
А вот первые две строки DataFrame (далее — датафрейм), data
.
user | question | score |
e213cc2b-387e-4d7d-983c-8abc19a586b1 | d3bdb068-7245-4521-ae57-d0e9692cb627 | 1 |
951ffaee-6e17-4599-a8c0-9dfd00470cd9 | d3bdb068-7245-4521-ae57-d0e9692cb627 | 0 |
Для проверки корректности нашего оптимизированного кода, мы можем использовать вывод исходного кода.
Поскольку мы пытаемся оптимизировать внутренний цикл, поместим его в собственную функцию, чтобы профилировать с помощью line_profiler.
Avg time per iteration: 35 ms
Speedup over baseline: 1.0x
% Time Line Contents
=====================
def compute_corrs(
qs_iter: Iterable, q_to_score: pd.DataFrame, grand_totals: pd.DataFrame
):
0.0 result = []
0.0 for qs in qs_iter:
13.5 qs_data = q_to_score.loc[islice[qs, :], :].swaplevel()
70.1 answered_all = qs_data.groupby(level=[0]).size() == K
0.4 answered_all = answered_all[answered_all].index
0.0 qs_total = (
6.7 qs_data.loc[islice[answered_all, :]]
1.1 .groupby(level=[0])
0.6 .sum()
0.3 .rename(columns={"score": "qs"})
)
7.4 r = qs_total.join(grand_totals).corr().qs.grand_total
0.0 result.append({"qs": qs, "r": r})
0.0 return result
Мы видим значения, которые пытаемся оптимизировать (среднее время итерации/ускорение), а также долю времени, потраченного на выполнение каждой строки.
Это позволяет оптимизировать код следующим образом:
Запускаем профилировщик
Определяем самые медленные строки
Пробуем сделать медленные строки более быстрыми
Повторяем
В приведённом выше коде мы видим, что есть наиболее медленная строка, которая занимает ~70% времени.
Однако есть еще один важный шаг, который предшествует вышеупомянутым:
Проверяем вывод на корректность
Запускаем профилировщик
Определяем самые медленные строки
Пробуем сделать медленные строки более быстрыми
Повторяем
Проверки корректности вывода помогают экспериментировать, пробовать различные методы, библиотеки и т.д., зная при этом, что любые случайные изменения в вычисляемой информации будут отслежены.
Оптимизация 1. Словарь множеств пользователей, ответивших на вопросы: users_who_answered_q
Наш базовый код выполняет различные тяжёлые операции Pandas, выясняя, какие пользователи ответили на заданный набор вопросов — qs
. В частности, для этого он проверяет каждую строку датафрейма, чтобы определить, какие пользователи отвечали на вопросы. В первой оптимизации вместо полноценного датафрейма мы можем использовать словарь множеств пользователей. Это позволит нам быстро выяснить, какие пользователи ответили на каждый вопрос qs
, и использовать пересечение множеств в Python, чтобы выявить пользователей, ответивших на все вопросы.
Avg time per iteration: 10.0 ms
Speedup over baseline: 3.5x
% Time Line Contents
=====================
def compute_corrs(qs_iter, users_who_answered_q, q_to_score, grand_totals):
0.0 result = []
0.0 for qs in qs_iter:
0.0 user_sets_for_qs = [users_who_answered_q[q] for q in qs]
3.6 answered_all = set.intersection(*user_sets_for_qs)
40.8 qs_data = q_to_score.loc[islice[qs, :], :].swaplevel()
0.0 qs_total = (
22.1 qs_data.loc[islice[list(answered_all), :]]
3.7 .groupby(level=[0])
1.9 .sum()
1.1 .rename(columns={"score": "qs"})
)
26.8 r = qs_total.join(grand_totals).corr().qs.grand_total
0.0 result.append({"qs": qs, "r": r})
0.0 return result
Так мы значительно ускоряем вычисление строки answered_all
, которая вместо 70 % теперь занимает 4 %, и наш код становится быстрее в 3 раза.
Оптимизация 2. Словарь score_dict
Если сложить время, затрачиваемое на каждую строку, участвующую в вычислении qs_total
(включая строку qs_data
), то получится ~65%; таким образом, наша следующая задача по оптимизации ясна. Нужно снова заменить тяжёлые операции над полным датафреймом (индексирование, группировка и т. д.) быстрым поиском по словарю. Для этого вводим score_dict
, словарь, который позволяет проводить оценку для заданной пары вопрос-пользователь.
Avg time per iteration: 690 μs
Speedup over baseline: 50.8x
% Time Line Contents
=====================
def compute_corrs(qs_iter, users_who_answered_q, score_dict, grand_totals):
0.0 result = []
0.0 for qs in qs_iter:
0.1 user_sets_for_qs = [users_who_answered_q[q] for q in qs]
35.9 answered_all = set.intersection(*user_sets_for_qs)
3.4 qs_total = {u: sum(score_dict[q, u] for q in qs) for u in answered_all}
8.6 qs_total = pd.DataFrame.from_dict(qs_total, orient="index", columns=["qs"])
0.1 qs_total.index.name = "user"
51.8 r = qs_total.join(grand_totals).corr().qs.grand_total
0.0 result.append({"qs": qs, "r": r})
0.0 return result
Это помогает нам ускорить код в 50 раз.
Оптимизация 3. Словарь grand_totals и np.corrcoef
Самая медленная строка в коде выше делает несколько вещей: сперва Pandas join
, чтобы объединить grand_totals
и qs_total
, а затем вычисляет для этого коэффициент корреляции. Опять же, мы можем ускорить процесс, используя поиск по словарю вместо join, и поскольку у нас больше нет объектов Pandas, используем np.corrcoef
вместо Pandas corr
.
Avg time per iteration: 380 μs
Speedup over baseline: 91.6x
% Time Line Contents
=====================
def compute_corrs(qs_iter, users_who_answered_q, score_dict, grand_totals):
0.0 result = []
0.0 for qs in qs_iter:
0.2 user_sets_for_qs = [users_who_answered_q[q] for q in qs]
83.9 answered_all = set.intersection(*user_sets_for_qs)
7.2 qs_total = [sum(score_dict[q, u] for q in qs) for u in answered_all]
0.5 user_grand_total = [grand_totals[u] for u in answered_all]
8.1 r = np.corrcoef(qs_total, user_grand_total)[0, 1]
0.1 result.append({"qs": qs, "r": r})
0.0 return result
Получаем ~90-кратное ускорение кода.
Оптимизация 4. Преобразование строк uuid в ints
Эта оптимизация не вносит никаких изменений в код внутреннего цикла. Но она ускоряет некоторые операции. Мы заменяем длинные uuid пользователя/вопроса (например, e213cc2b-387e-4d7d-983c-8abc19a586b1) на гораздо более короткие целочисленные данные. Как это делается:
data.user = data.user.map({u: i for i, u in enumerate(data.user.unique())})
data.question = data.question.map(
{q: i for i, q in enumerate(data.question.unique())}
)
Теперь измеряем:
Avg time per iteration: 210 μs
Speedup over baseline: 168.5x
% Time Line Contents
=====================
def compute_corrs(qs_iter, users_who_answered_q, score_dict, grand_totals):
0.0 result = []
0.1 for qs in qs_iter:
0.4 user_sets_for_qs = [users_who_answered_q[q] for q in qs]
71.6 answered_all = set.intersection(*user_sets_for_qs)
13.1 qs_total = [sum(score_dict[q, u] for q in qs) for u in answered_all]
0.9 user_grand_total = [grand_totals[u] for u in answered_all]
13.9 r = np.corrcoef(qs_total, user_grand_total)[0, 1]
0.1 result.append({"qs": qs, "r": r})
0.0 return result
Оптимизация 5. Массив np.bool_ вместо множеств пользователей
Видно, что операция с множествами пользователей в коде выше по-прежнему самая медленная. Вместо использования наборов ints мы переходим к использованию массива пользователей np.bool_
и применяем np.logical_and.reduce
для поиска пользователей, ответивших на все вопросы qs. (Обратите внимание, что np.bool_
использует целый байт для каждого элемента, но np.logical_and.reduce
все равно довольно быстр.) Это даёт нам значительное ускорение:
Avg time per iteration: 75 μs
Speedup over baseline: 466.7x
% Time Line Contents
=====================
def compute_corrs(qs_iter, users_who_answered_q, score_dict, grand_totals):
0.0 result = []
0.1 for qs in qs_iter:
12.0 user_sets_for_qs = users_who_answered_q[qs, :] # numpy indexing
9.9 answered_all = np.logical_and.reduce(user_sets_for_qs)
10.7 answered_all = np.where(answered_all)[0]
33.7 qs_total = [sum(score_dict[q, u] for q in qs) for u in answered_all]
2.6 user_grand_total = [grand_totals[u] for u in answered_all]
30.6 r = np.corrcoef(qs_total, user_grand_total)[0, 1]
0.2 result.append({"qs": qs, "r": r})
0.0 return result
Оптимизация 6. score_matrix вместо словаря
Теперь самая медленная строка — вычисление qs_total
. Следуя примеру из оригинальной статьи, мы переходим к использованию плотного массива np.array для поиска оценок вместо словаря, и используем быструю индексацию NumPy для получения оценок.
Avg time per iteration: 56 μs
Speedup over baseline: 623.7x
% Time Line Contents
=====================
def compute_corrs(qs_iter, users_who_answered_q, score_matrix, grand_totals):
0.0 result = []
0.2 for qs in qs_iter:
16.6 user_sets_for_qs = users_who_answered_q[qs, :]
14.0 answered_all = np.logical_and.reduce(user_sets_for_qs)
14.6 answered_all = np.where(answered_all)[0]
7.6 qs_total = score_matrix[answered_all, :][:, qs].sum(axis=1)
3.9 user_grand_total = [grand_totals[u] for u in answered_all]
42.7 r = np.corrcoef(qs_total, user_grand_total)[0, 1]
0.4 result.append({"qs": qs, "r": r})
0.0 return result
Оптимизация 7. Реализация corrcoef
Самая медленная строка в коде — np.corrcoef
... Мы всеми силами пытаемся оптимизировать код, поэтому вот наша собственная реализация corrcoef, которая в данном случае будет в два раза быстрее:
def corrcoef(a: list[float], b: list[float]) -> float | None:
"""same as np.corrcoef(a, b)[0, 1]"""
n = len(a)
sum_a = sum(a)
sum_b = sum(b)
sum_ab = sum(a_i * b_i for a_i, b_i in zip(a, b))
sum_a_sq = sum(a_i**2 for a_i in a)
sum_b_sq = sum(b_i**2 for b_i in b)
num = n * sum_ab - sum_a * sum_b
den = sqrt(n * sum_a_sq - sum_a**2) * sqrt(n * sum_b_sq - sum_b**2)
if den == 0:
return None
return num / den
Получаем приличное ускорение:
Avg time per iteration: 43 μs
Speedup over baseline: 814.6x
% Time Line Contents
=====================
def compute_corrs(qs_iter, users_who_answered_q, score_matrix, grand_totals):
0.0 result = []
0.2 for qs in qs_iter:
21.5 user_sets_for_qs = users_who_answered_q[qs, :] # numpy indexing
18.7 answered_all = np.logical_and.reduce(user_sets_for_qs)
19.7 answered_all = np.where(answered_all)[0]
10.0 qs_total = score_matrix[answered_all, :][:, qs].sum(axis=1)
5.3 user_grand_total = [grand_totals[u] for u in answered_all]
24.1 r = corrcoef(qs_total, user_grand_total)
0.5 result.append({"qs": qs, "r": r})
0.0 return result
Оптимизация 8. Преждевременное внедрение Numba
Мы ещё не закончили оптимизацию структур данных в приведённом выше коде, но давайте посмотрим, что нам даст внедрение на текущем этапе Numba? Речь о библиотеке в экосистеме Python, которая «переводит подмножество кода Python и NumPy в быстрый машинный код».
Чтобы иметь возможность использовать Numba, выполним два изменения:
Модификация 1. Передаем qs_combinations как массив numpy, вместо qs_iter
Numba не очень хорошо работает с itertools
или генераторами, поэтому мы заранее превращаем qs_iter
в массив NumPy, чтобы передать его функции. Влияние этого изменения на скорость выполнения (до добавления Numba) показано ниже.
Avg time per iteration: 42 μs
Speedup over baseline: 829.2x
Модификация 2. Массив результатов вместо списка
Вместо добавления в список, мы инициализируем массив и помещаем в него результаты. Вот как это изменение повлияло на скорость.
Avg time per iteration: 42 μs
Speedup over baseline: 833.8x
В итоге наш код выглядит так:
import numba
@numba.njit(parallel=False)
def compute_corrs(qs_combinations, users_who_answered_q, score_matrix, grand_totals):
result = np.empty(len(qs_combinations), dtype=np.float64)
for i in numba.prange(len(qs_combinations)):
qs = qs_combinations[i]
user_sets_for_qs = users_who_answered_q[qs, :]
# numba doesn't support np.logical_and.reduce
answered_all = user_sets_for_qs[0]
for j in range(1, len(user_sets_for_qs)):
answered_all *= user_sets_for_qs[j]
answered_all = np.where(answered_all)[0]
qs_total = score_matrix[answered_all, :][:, qs].sum(axis=1)
user_grand_total = grand_totals[answered_all]
result[i] = corrcoef_numba(qs_total, user_grand_total)
return result
(Обратите внимание, что мы также дополнили corrcoef
с помощью Numba, потому что функции, вызываемые внутри функции Numba, тоже должны быть скомпилированы.)
Результаты с параметром parallel=False
:
Avg time per iteration: 47 μs
Speedup over baseline: 742.2x
Результаты с параметром parallel=True
:
Avg time per iteration: 8.5 μs
Speedup over baseline: 4142.0x
Видно, что при значении parallel=False
код Numba работает немного медленнее, чем наш предыдущий код на Python, но когда мы включаем параллелизм, то начинаем использовать все ядра процессора (10 на нашей рабочей машине) — и это даёт хороший множитель скорости.
Однако мы теряем возможность использовать line_profiler на JIT-компилированном коде; (возможно, мы захотим обратиться к сгенерированному LLVM IR / сборке).
Оптимизация 9. Bitsets, без Numba
Пока отвлечёмся от Numba. В оригинальной статье для быстрого вычисления пользователей, ответивших на текущий qs
, используются bitsets — проверим, применим ли такой подход в нашем случае. Для реализации bitsets мы можем использовать массивы NumPy np.int64
и np.bitwise_and.reduce
. В отличие от использования массива np.bool_
, теперь мы используем отдельные биты в байте для представления сущностей в множестве. Обратите внимание, что для данного bitset может понадобиться несколько байтов, в зависимости от максимального количества элементов, которые нам нужны. Мы можем использовать быстрый bitwise_and для байтов каждого вопроса в qs
, чтобы найти пересечение множеств и, следовательно, количество пользователей, ответивших на все qs
.
Вот функции bitset
, которые мы будем использовать:
def bitset_create(size):
"""Initialise an empty bitset"""
size_in_int64 = int(np.ceil(size / 64))
return np.zeros(size_in_int64, dtype=np.int64)
def bitset_add(arr, pos):
"""Add an element to a bitset"""
int64_idx = pos // 64
pos_in_int64 = pos % 64
arr[int64_idx] |= np.int64(1) << np.int64(pos_in_int64)
def bitset_to_list(arr):
"""Convert a bitset back into a list of ints"""
result = []
for idx in range(arr.shape[0]):
if arr[idx] == 0:
continue
for pos in range(64):
if (arr[idx] & (np.int64(1) << np.int64(pos))) != 0:
result.append(idx * 64 + pos)
return np.array(result)
И мы можем инициализировать bitsets следующим образом:
users_who_answered_q = np.array(
[bitset_create(data.user.nunique()) for _ in range(data.question.nunique())]
)
for q, u in data[["question", "user"]].values:
bitset_add(users_who_answered_q[q], u)
Посмотрим, какое ускорение мы получим:
Avg time per iteration: 550 μs
Speedup over baseline: 64.2x
% Time Line Contents
=====================
def compute_corrs(qs_combinations, users_who_answered_q, score_matrix, grand_totals):
0.0 num_qs = qs_combinations.shape[0]
0.0 bitset_size = users_who_answered_q[0].shape[0]
0.0 result = np.empty(qs_combinations.shape[0], dtype=np.float64)
0.0 for i in range(num_qs):
0.0 qs = qs_combinations[i]
0.3 user_sets_for_qs = users_who_answered_q[qs_combinations[i]]
0.4 answered_all = np.bitwise_and.reduce(user_sets_for_qs)
96.7 answered_all = bitset_to_list(answered_all)
0.6 qs_total = score_matrix[answered_all, :][:, qs].sum(axis=1)
0.0 user_grand_total = grand_totals[answered_all]
1.9 result[i] = corrcoef(qs_total, user_grand_total)
0.0 return result
Как видно, мы замедлились, поскольку операция bitset_to_list
занимает слишком много времени.
Оптимизация 10. Numba на bitset_to_list
Преобразуем bitset_to_list
в скомпилированный код. Для этого мы можем добавить декоратор Numba:
@numba.njit
def bitset_to_list(arr):
result = []
for idx in range(arr.shape[0]):
if arr[idx] == 0:
continue
for pos in range(64):
if (arr[idx] & (np.int64(1) << np.int64(pos))) != 0:
result.append(idx * 64 + pos)
return np.array(result)
Измерим скорость:
Benchmark #14: bitsets, with numba on bitset_to_list
Using 1000 iterations...
Avg time per iteration: 19 μs
Speedup over baseline: 1801.2x
% Time Line Contents
=====================
def compute_corrs(qs_combinations, users_who_answered_q, score_matrix, grand_totals):
0.0 num_qs = qs_combinations.shape[0]
0.0 bitset_size = users_who_answered_q[0].shape[0]
0.0 result = np.empty(qs_combinations.shape[0], dtype=np.float64)
0.3 for i in range(num_qs):
0.6 qs = qs_combinations[i]
8.1 user_sets_for_qs = users_who_answered_q[qs_combinations[i]]
11.8 answered_all = np.bitwise_and.reduce(user_sets_for_qs)
7.7 answered_all = bitset_to_list(answered_all)
16.2 qs_total = score_matrix[answered_all, :][:, qs].sum(axis=1)
1.1 user_grand_total = grand_totals[answered_all]
54.1 result[i] = corrcoef(qs_total, user_grand_total)
0.0 return result
Мы получили ускорение в 1800 раз по сравнению с исходным кодом. Вспомните, что оптимизация 7, до введения Numba, дала 814x. (Оптимизация 8 дала 4142x, но это было с parallel=True
во внутреннем цикле, так что показатель здесь нерелевантен.)
Оптимизация 11. Numba на corrcoef
Строчка с corrcoef снова выделяется как слишком медленная. Навесим на corrcoef
декоратор Numba.
@numba.njit
def corrcoef_numba(a, b):
"""same as np.corrcoef(a, b)[0, 1]"""
n = len(a)
sum_a = sum(a)
sum_b = sum(b)
sum_ab = sum(a * b)
sum_a_sq = sum(a * a)
sum_b_sq = sum(b * b)
num = n * sum_ab - sum_a * sum_b
den = math.sqrt(n * sum_a_sq - sum_a**2) * math.sqrt(n * sum_b_sq - sum_b**2)
return np.nan if den == 0 else num / den
Смотрим результаты:
Avg time per iteration: 11 μs
Speedup over baseline: 3218.9x
% Time Line Contents
=====================
def compute_corrs(qs_combinations, users_who_answered_q, score_matrix, grand_totals):
0.0 num_qs = qs_combinations.shape[0]
0.0 bitset_size = users_who_answered_q[0].shape[0]
0.0 result = np.empty(qs_combinations.shape[0], dtype=np.float64)
0.7 for i in range(num_qs):
1.5 qs = qs_combinations[i]
15.9 user_sets_for_qs = users_who_answered_q[qs_combinations[i]]
26.1 answered_all = np.bitwise_and.reduce(user_sets_for_qs)
16.1 answered_all = bitset_to_list(answered_all)
33.3 qs_total = score_matrix[answered_all, :][:, qs].sum(axis=1)
2.0 user_grand_total = grand_totals[answered_all]
4.5 result[i] = corrcoef_numba(qs_total, user_grand_total)
0.0 return result
Прекрасно, очередное значительное ускорение!
Оптимизация 12. Numba на bitset_and
Вместо использования np.bitwise_and.reduce
мы вводим функцию bitwise_and
и применяем к ней jit-компиляцию.
@numba.njit
def bitset_and(arrays):
result = arrays[0].copy()
for i in range(1, len(arrays)):
result &= arrays[i]
return result
Benchmark #16: numba also on bitset_and
Using 1000 iterations...
Avg time per iteration: 8.9 μs
Speedup over baseline: 3956.7x
% Time Line Contents
=====================
def compute_corrs(qs_combinations, users_who_answered_q, score_matrix, grand_totals):
0.1 num_qs = qs_combinations.shape[0]
0.0 bitset_size = users_who_answered_q[0].shape[0]
0.1 result = np.empty(qs_combinations.shape[0], dtype=np.float64)
1.0 for i in range(num_qs):
1.5 qs = qs_combinations[i]
18.4 user_sets_for_qs = users_who_answered_q[qs_combinations[i]]
16.1 answered_all = bitset_and(user_sets_for_qs)
17.9 answered_all = bitset_to_list(answered_all)
37.8 qs_total = score_matrix[answered_all, :][:, qs].sum(axis=1)
2.4 user_grand_total = grand_totals[answered_all]
4.8 result[i] = corrcoef_numba(qs_total, user_grand_total)
0.0 return result
Оптимизация 13. Numba на всю функцию
Код стал значительно быстрее исходного, причём вычисления распределены довольно равномерно между несколькими строками цикла. Похоже, самая медленная строка выполняет индексацию NumPy, которая и так довольно быстрая. Давайте скомпилируем всю функцию с помощью Numba.
@numba.njit(parallel=False)
def compute_corrs(qs_combinations, users_who_answered_q, score_matrix, grand_totals):
result = np.empty(len(qs_combinations), dtype=np.float64)
for i in numba.prange(len(qs_combinations)):
qs = qs_combinations[i]
user_sets_for_qs = users_who_answered_q[qs, :]
answered_all = user_sets_for_qs[0]
# numba doesn't support np.logical_and.reduce
for j in range(1, len(user_sets_for_qs)):
answered_all *= user_sets_for_qs[j]
answered_all = np.where(answered_all)[0]
qs_total = score_matrix[answered_all, :][:, qs].sum(axis=1)
user_grand_total = grand_totals[answered_all]
result[i] = corrcoef_numba(qs_total, user_grand_total)
return result
Avg time per iteration: 4.2 μs
Speedup over baseline: 8353.2x
А теперь с параметром parallel=True
:
Avg time per iteration: 960 ns
Speedup over baseline: 36721.4x
Отлично, наш код уже в 36 000 раз быстрее исходного.
Оптимизация 14. Numba, встраивание с накоплением вместо массивов
Куда двигаться дальше?... Ну, в коде все ещё достаточно много значений помещается в массивы, а затем передаётся по ним. Если мы посмотрим, как вычисляется corrcoef, то поймём, что нам не нужно создавать массивы answered_all
и user_grand_total
, мы можем накапливать значения по мере выполнения цикла.
Вот код (мы также включили некоторые оптимизации компилятора, например, отключили boundschecking
для массивов и включили fastmath
):
@numba.njit(boundscheck=False, fastmath=True, parallel=False, nogil=True)
def compute_corrs(qs_combinations, users_who_answered_q, score_matrix, grand_totals):
num_qs = qs_combinations.shape[0]
bitset_size = users_who_answered_q[0].shape[0]
corrs = np.empty(qs_combinations.shape[0], dtype=np.float64)
for i in numba.prange(num_qs):
# bitset will contain users who answered all questions in qs_array[i]
bitset = users_who_answered_q[qs_combinations[i, 0]].copy()
for q in qs_combinations[i, 1:]:
bitset &= users_who_answered_q[q]
# retrieve stats for the users to compute correlation
n = 0.0
sum_a = 0.0
sum_b = 0.0
sum_ab = 0.0
sum_a_sq = 0.0
sum_b_sq = 0.0
for idx in range(bitset_size):
if bitset[idx] != 0:
for pos in range(64):
if (bitset[idx] & (np.int64(1) << np.int64(pos))) != 0:
user_idx = idx * 64 + pos
score_for_qs = 0.0
for q in qs_combinations[i]:
score_for_qs += score_matrix[user_idx, q]
score_for_user = grand_totals[user_idx]
n += 1.0
sum_a += score_for_qs
sum_b += score_for_user
sum_ab += score_for_qs * score_for_user
sum_a_sq += score_for_qs * score_for_qs
sum_b_sq += score_for_user * score_for_user
num = n * sum_ab - sum_a * sum_b
den = np.sqrt(n * sum_a_sq - sum_a**2) * np.sqrt(n * sum_b_sq - sum_b**2)
corrs[i] = np.nan if den == 0 else num / den
return corrs
Посмотрим со значением parallel=False
.
Avg time per iteration: 1.7 μs
Speedup over baseline: 20850.5x
Результат можно сравнить с оптимизацией 12 с parallel=False
, которая показала 8353x.
Теперь с параметром parallel=True
.
Avg time per iteration: 210 ns
Speedup over baseline: 170476.3x
Мы достигли ускорения в 170 000 по сравнению с исходным кодом!
Вывод
Благодаря Numba и NumPy мы получили большинство тех инструментов, которые сделали оптимизированный код Rust быстрым: в частности, bitsets, SIMD и параллелизм на уровне циклов. Сперва мы значительно ускорили оригинальный код на Python с помощью нескольких вспомогательных функций с JIT-компиляцией, в итоге использовали JIT-компиляцию повсеместно и оптимизировали код для этого. Мы использовали подход проб и ошибок, применяя профилирование, чтобы сосредоточить усилия на самых медленных строках кода. Мы показали, что можем использовать Numba для постепенного добавления JIT-компилированного кода в нашу кодовую базу Python. Мы можем сразу же добавить этот код в существующую кодовую базу Python. Однако мы не достигли 180 000-кратного ускорения оптимизированного кода Rust, и развернули собственную реализацию корреляции и bitsets, в то время как код Rust смог использовать библиотеки для них, оставаясь при этом быстрым.
Это было забавное упражнение, которое, надеюсь, продемонстрировало вам некоторые полезные инструменты в экосистеме Python.
Стал бы я рекомендовать один подход вместо другого? Нет, все зависит от конкретной ситуации.