numba,一个强大的 python 库
大家好,我是小寒。
今天给大家分享一个强大的 python 库,numba。
https://github.com/numba/numba
Numba 可用于加速 Python 函数,而无需编写任何 C 或 C++ 代码。
Numba 的安装
pip install numba
如何对 Pandas 使用 Numba?
将 Numba 与 Pandas 结合使用有两种方法:
在 pandas 方法中指定 engine="numba" 关键字。 定义你自己的用 @jit 修饰的 Python 函数,并将 DataFrame 或 Series 的底层 NumPy 数组(使用 to_numpy())传递到函数中。
使用 engine="numba"
你可以通过指定 engine="numba" 来增强所选 Pandas 方法的执行。
滚动平均值、中位数、最大值、最小值、总和和标准差等。 groupby 滚动平均值、中位数、最大值、最小值、总和和标准差等。 expanding 平均值、中位数、最大值、最小值、总和和标准差等。 groupby expanding 平均值、中位数、最大值、最小值、总和和标准差等。
# create a sample DataFrame
data = np.random.rand(int(1e5), 4)
df = pd.DataFrame(data)
window_size = 10
# rolling sum with Numba
rolling_sum = df.rolling(window_size).sum(engine='numba')
定义你自己的函数
你还可以定义用 @jit 或 @njit 装饰的 Python 函数,并将 DataFrame 或 Series 的基础 NumPy 数组(使用 to_numpy())传递到函数中。此方法比使用 engine="numba" 更灵活,因为你可以定义自己的自定义函数。
from numba import njit
# create a sample DataFrame
data = np.random.rand(int(1e5), 4)
df = pd.DataFrame(data)
# Define the custom function
@njit
def sum_of_squares(x):
return np.sum(x**2, axis=0)
result = sum_of_squares(df.to_numpy())
它采用 NumPy 数组作为输入并计算沿列的平方和。
性能比较
import time
import matplotlib.pyplot as plt
import numba
import numpy as np
import pandas as pd
plt.style.use('ggplot')
# Define the custom rolling apply function
def root_mean_square(x):
return np.sqrt(np.mean(x**2))
def bench(df: pd.DataFrame, use_numba: bool, use_parallel: bool = False, n_times: int = 10):
engine = 'numba' if use_numba else None
engine_kwargs = {'parallel': True, 'nopython': True} if use_parallel else None
elapsed_time_list = []
for _ in range(n_times):
start_time = time.time()
df.rolling(window_size).apply(root_mean_square, raw=True, engine=engine, engine_kwargs=engine_kwargs)
elapsed_time_list.append(time.time() - start_time)
return np.mean(elapsed_time_list)
# Generate sample data
np.random.seed(42)
data = np.random.rand(int(1e5), 4)
df = pd.DataFrame(data)
window_sizes = range(10, 1001, 100)
mean_time_without_numba = []
mean_time_with_numba = []
mean_time_with_numba_parallel = []
# Benchmark the performance for different window sizes
for window_size in window_sizes:
# Without Numba
mean_time = bench(df, use_numba=False)
mean_time_without_numba.append(mean_time)
mean_time = bench(df, use_numba=True)
mean_time_with_numba.append(mean_time)
mean_time = bench(df, use_numba=True, use_parallel=True)
mean_time_with_numba_parallel.append(mean_time)
mean_time_without_numba = np.array(mean_time_without_numba)
mean_time_with_numba = np.array(mean_time_with_numba)
mean_time_with_numba_parallel = np.array(mean_time_with_numba_parallel)
print("mean_time_without_numba", mean_time_without_numba)
print("mean_time_with_numba", mean_time_with_numba)
print("mean_time_with_numba_parallel", mean_time_with_numba_parallel)
fig, ax = plt.subplots(figsize=(8, 6))
ax.plot(window_sizes, mean_time_without_numba / mean_time_with_numba, label='Numba')
ax.plot(window_sizes, mean_time_without_numba / mean_time_with_numba_parallel, label='Numba + Parallel')
ax.set_xlabel('Window Size')
ax.set_ylabel('Speedup')
ax.set_title('Speedup compared to without Numba: Pandas Rolling Apply')
ax.legend()
fig.savefig("numba_speedup_pandas_rolling_apply.png")
如上图所示,Numba 将 pandas 滚动函数的执行速度提高了 260 倍。
与没有 numba 的非优化版本相比,Numba 显着提高了 pandas 滚动的性能。平均执行时间减少了一个数量级。
与没有并行化的 numba 相比,具有并行化的 Numba 进一步增强了性能,实现了额外的加速。
随着窗口大小的增加,加速比降低,这可能是由于所需滚动操作数量的减少。
定义 Numba 修饰的函数可以灵活地优化自定义操作。
最后
—
「进群方式:加我微信,备注 “python”」
往期回顾
Fashion-MNIST 服装图片分类-Pytorch实现
如果对本文有疑问可以加作者微信直接交流。进技术交流群的可以加微信拉你进群。