Numbaを使用して次のコードを高速化する

Question

Numbaを使用してコードを高速化しようとしています。コードはシンプルですが、基本的にループは単純な数値計算配列で計算されます。

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を使用して最適化し、定義してみました。私はNumbaははるかに高速であると予想しながら

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

しかし、時間は、ほぼ同じです。

テストコード：

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

ジッタ関数の実行を1回だけ実行すると、実行時間とNumbaがコードをジットするのにかかる時間の両方が表示されます。 2回目にコードを実行すると、Numbaはコンパイルされた関数のメモリ内キャッシュを使用するので、引数の種類ごとにコンパイル時間だけを支払うので、実際のスピードアップが表示されます。

私のマシンで、Python 3.6とnumba 0.31.0を使用している場合、純粋なpython関数は0.32秒かかります。私が初めてfastMinimizeSquareDiffBudgetと呼ぶときは0.57秒かかるが、2回目は0.31秒かかる。

Numbaがnopythonモードでコンパイルできない関数を持っているので、大きなスピードアップが見られない理由は、はるかに遅くobject modeに落ちます。 nopython=Trueをjitメソッドに渡すと、コンパイルできない箇所が表示されます。私が見た2つの問題は、np.size(x,0)の代わりにx.shape[0]を使用する必要があり、あなたの方法でminを使用できないということでした。