使用 Numba 加速以下代码

Question

我正在尝试使用 Numba 来加速一段代码。代码很简单，基本上是一个对 numpy 数组进行简单计算的循环。

import numpy as np
import time
from numba import jit, double

def MinimizeSquareDiffBudget(x, budget):
    if (budget > np.sum(x)):
        return x
    n = np.size(x,0)
    j = 1
    i = 0
    y = np.zeros((n, 1))
    while (budget > 0):
        while (x[i] == x[j]) and (j < n-1):
            j += 1
        i = j - 1
        if (np.std(x)<1e-10):
            to_give = budget/n
            y += to_give
            x= x- to_give
            break
        to_give = min(budget, (x[0] - x[j])*j)
        y[0:j] += to_give/j
        x[0:j]=x[0:j]-to_give/j
        budget = budget - to_give
        j = 1
    return y

现在，我尝试使用@jit 并通过定义对其进行优化：

fastMinimizeSquareDiffBudget = jit(double[:,:](double[:,:], double[:,:]))(MinimizeSquareDiffBudget)

但是，时间大致相同，而我预计 Numba 会快得多。

测试代码：

budget = 335.0

x = np.random.uniform(0,1,(1000,1))
x.sort(axis=0)
x = x[::-1]
t = time.process_time()
y = MinimizeSquareDiffBudget(x, budget)
print(time.process_time()-t)

x = np.random.uniform(0,1,(1000,1))
x.sort(axis=0)
x = x[::-1]

t = time.process_time()
y = fastMinimizeSquareDiffBudget(x, budget)
print(time.process_time()-t)

直接实现需要 0.28 秒，使用 Numba 优化代码需要 0.45 秒。用 C 编写相同的代码只需要不到 0.001 秒。

有什么想法吗？

Answer 1

当您只对 jitted 函数的一次执行计时时，您会同时看到运行时间和 Numba jit 代码所花费的时间。如果你第二次运行代码，你会看到实际的加速，因为 Numba 使用编译函数的内存缓存，所以你只需为每个参数类型支付一次编译时间。

在我使用 python 3.6 和 numba 0.31.0 的机器上，纯 python 函数需要 0.32 秒。我第一次调用 fastMinimizeSquareDiffBudget 需要 0.57 秒，但第二次需要 0.31 秒。

现在你没有看到巨大加速的原因是因为你有一个 Numba 无法在 nopython 模式下编译的函数，所以它回落到更慢的 object mode.如果将 nopython=True 传递给 jit 方法，您将能够看到它无法编译的地方。我看到的两个问题是你应该使用 x.shape[0] 而不是 np.size(x,0)，你不能像现在这样使用 min。

使用 Numba 加速以下代码

Accelerating the following code using Numba

python

numpy

numba