查看原文
其他

numba,一个强大的 python 库

程序员小寒 程序员学长 2023-09-17

大家好,我是小寒。

今天给大家分享一个强大的 python 库,numba

https://github.com/numba/numba

Pandas 是 Python 中流行的数据分析库。然而,随着数据集大小的增长,原生 Python 代码对于滚动窗口计算等关键操作可能会变得很慢。这就是 Numba 用武之地。
Numba 是 Python 的即时编译器,可在运行时将 Python 代码转换为机器代码

Numba 可用于加速 Python 函数,而无需编写任何 C 或 C++ 代码。

在本文中,我们将探讨如何利用 Numba 加速常见的 Pandas 工作流程,例如滚动统计。‍
你将了解启用 Numba 编译所需的简单注释以及如何使用并行执行等选项来调整性能。

Numba 的安装

我们可以直接使用pip 来进行安装。
pip install numba

如何对 Pandas 使用 Numba?

将 Numba 与 Pandas 结合使用有两种方法:

  • 在 pandas 方法中指定 engine="numba" 关键字
  • 定义你自己的用 @jit 修饰的 Python 函数,并将 DataFrame 或 Series 的底层 NumPy 数组(使用 to_numpy())传递到函数中。

使用 engine="numba"

你可以通过指定 engine="numba" 来增强所选 Pandas 方法的执行。

这指示该方法利用 Numba 来加速性能。支持该关键字的方法主要涉及窗口操作,例如:
  • 滚动平均值、中位数、最大值、最小值、总和和标准差等。
  • 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')
上面的代码片段生成一个随机 DataFrame 并使用 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())
此代码片段使用 Numba 的 @njit 装饰器(即 @jit(nopython=True) 的别名)来定义名为 sum_of_squares 的自定义函数。

它采用 NumPy 数组作为输入并计算沿列的平方和。

该代码生成一个随机 DataFrame,使用 to_numpy() 将其转换为 NumPy 数组,并应用 sum_of_squares 函数计算每列的平方和。结果存储在 result 变量中。

性能比较

我们比较执行时间并评估通过利用 Numba 的即时编译功能获得的潜在加速。
为此,我们在如下三种场景下测量代码的执行时间:有和没有 Numba 优化,以及并行执行。
下面提供了执行这些基准测试的代码:
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 在 Pandas 中进行滚动操作所实现的显着性能改进。
  • 与没有 numba 的非优化版本相比,Numba 显着提高了 pandas 滚动的性能。平均执行时间减少了一个数量级。

  • 与没有并行化的 numba 相比,具有并行化的 Numba 进一步增强了性能,实现了额外的加速。

  • 随着窗口大小的增加,加速比降低,这可能是由于所需滚动操作数量的减少。

Numba 即时编译器可以通过为关键函数生成优化的机器代码,将 Pandas 操作加速高达 260 倍。在 Pandas 方法中使用 numba 引擎可以轻松地在列上进行多线程处理。

定义 Numba 修饰的函数可以灵活地优化自定义操作。

尽管存在编译开销等限制,Numba 独特地将 Python 和 Pandas 的生产力与类似 C 的速度结合在一起。
对于处理大型数据集的数据科学家来说,Numba 是加速 Pandas 中重复数值计算的不可或缺的工具。


最后



今天的分享就到这里。如果觉得不错,点赞,转发安排起来吧。接下来我们会分享更多的 「深度学习案例以及python相关的技术」,欢迎大家关注。最后,最近新建了一个 python 学习交流群,会经常分享 「python相关学习资料,也可以问问题,非常棒的一个群」

「进群方式:加我微信,备注 “python”」



往期回顾


Fashion-MNIST 服装图片分类-Pytorch实现

python 探索性数据分析(EDA)案例分享

深度学习案例分享 | 房价预测 - PyTorch 实现

万字长文 |  面试高频算法题之动态规划系列

面试高频算法题之回溯算法(全文六千字)  

    



如果对本文有疑问可以加作者微信直接交流。进技术交流群的可以加微信拉你进群。

您可能也对以下帖子感兴趣

文章有问题?点此查看未经处理的缓存