เร็วขึ้นสำหรับลูปด้วยอาร์เรย์ใน Python

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]

ฉันทำงานกับลูปนานมาก %%timeitด้านบนด้วยการpassแทนที่ผลตอบแทนการดำเนินการ

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

สิ่งนี้ไม่สามารถยอมรับได้ในบริบท (C ++ ใช้เวลา 6.5 วินาที) ไม่มีเหตุผลข้างต้นที่จะทำกับวัตถุ Python อาร์เรย์มีประเภทที่กำหนดไว้อย่างชัดเจน การนำสิ่งนี้ไปใช้ใน C / C ++ เป็นส่วนขยายถือเป็นการใช้งานมากเกินไปทั้งในส่วนของผู้พัฒนาและผู้ใช้ ฉันแค่ส่งอาร์เรย์เพื่อวนซ้ำและทำเลขคณิต

มีวิธีบอก Numpy ว่า "ย้ายตรรกะนี้ไปที่ C" หรือไม่หรือไลบรารีอื่นที่สามารถจัดการลูปที่ซ้อนกันซึ่งเกี่ยวข้องกับอาร์เรย์เท่านั้น ฉันขอเป็นกรณีทั่วไปไม่ใช่วิธีแก้ปัญหาสำหรับตัวอย่างเฉพาะนี้ (แต่ถ้าคุณมีฉันสามารถเปิดถาม & ตอบแยกต่างหากได้)

Answer 1

นี้นั้นเป็นความคิดที่อยู่เบื้องหลังNumba ไม่เร็วเท่า C แต่เข้าใกล้ได้ ... มันใช้คอมไพเลอร์ jit เพื่อคอมไพล์โค้ด python ไปยังเครื่องและเข้ากันได้กับฟังก์ชัน Numpy ส่วนใหญ่ (ในเอกสารคุณพบรายละเอียดทั้งหมด)

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

Python บริสุทธิ์:

%%timeit

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

338 ms ± 12.6 ms ต่อลูป (ค่าเฉลี่ย± std. dev ของการรัน 7 ครั้งแต่ละลูป 1 ครั้ง)

ด้วย Numba:

%%timeit

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

12 ms ± 534 s ต่อลูป (ค่าเฉลี่ย± std. dev ของ 7 รัน, 100 ลูปแต่ละครั้ง)