Python'da dizilerle daha hızlı döngüler

Question

N, M = 1000, 4000000
a = np.random.uniform(0, 1, (N, M))
k = np.random.randint(0, N, (N, M))

out = np.zeros((N, M))
for i in range(N):
    for j in range(M):
        out[k[i, j], j] += a[i, j]

Çok uzun döngülerle çalışıyorum; operasyon verimlerini değiştirerek %%timeityukarıdapass

1min 19s ± 663 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

bu bağlamda kabul edilemez (C ++ 6,5 saniye sürdü). Yukarıdakilerin Python nesneleriyle yapılması için bir neden yoktur; dizilerin iyi tanımlanmış türleri vardır. Bunu C / C ++ 'da bir uzantı olarak uygulamak, hem geliştirici hem de kullanıcı açısından aşırı bir şeydir; Sadece dizileri döngüye ve aritmetik işlemeye geçiriyorum.

Numpy'ye "bu mantığı C'ye taşı" demenin bir yolu var mı, yoksa yalnızca dizileri içeren iç içe döngüleri işleyebilecek başka bir kitaplık mı? Bunu genel durum için arıyorum, bu özel örnek için geçici çözümler değil (ancak eğer varsa, ayrı bir Soru-Cevap açabilirim).

Answer 1

Bu temelde Numba'nın arkasındaki fikirdir . C kadar hızlı değil, ancak yaklaşabilir ... Python kodunu makineye derlemek için bir jit derleyicisi kullanır ve çoğu Numpy işleviyle uyumludur. (Dokümanlarda tüm detayları bulursunuz)

import numpy as np
from numba import njit


@njit
def f(N, M):
    a = np.random.uniform(0, 1, (N, M))
    k = np.random.randint(0, N, (N, M))

    out = np.zeros((N, M))
    for i in range(N):
        for j in range(M):
            out[k[i, j], j] += a[i, j]
    return out


def f_python(N, M):
    a = np.random.uniform(0, 1, (N, M))
    k = np.random.randint(0, N, (N, M))

    out = np.zeros((N, M))
    for i in range(N):
        for j in range(M):
            out[k[i, j], j] += a[i, j]
    return out

Saf Python:

%%timeit

N, M = 100, 4000
f_python(M, N)

338 ms ± 12.6 ms döngü başına (ortalama ± std. Sapma 7 çalıştırma, her biri 1 döngü)

Numba ile:

%%timeit

N, M = 100, 4000
f(M, N)

Döngü başına 12 ms ± 534 µs (ortalama ± standart sapma 7 çalıştırma, her biri 100 döngü)