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运行代码所需的运行时间和时间。如果您第二次运行代码，您将看到实际的加速，因为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加速以下代码

1 个答案: