Vòng lặp for nhanh hơn với mảng trong 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]

Tôi làm việc với vòng lặp for rất dài; %%timeitở trên với việc passthay thế các kết quả hoạt động

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

điều này là không thể chấp nhận trong ngữ cảnh (C ++ mất 6,5 giây). Không có lý do gì ở trên được thực hiện với các đối tượng Python; mảng có các kiểu được xác định rõ ràng. Việc triển khai điều này trong C / C ++ dưới dạng tiện ích mở rộng là quá mức cần thiết đối với cả nhà phát triển và người dùng; Tôi chỉ truyền các mảng để lặp lại và tính toán số học.

Có cách nào để nói với Numpy "chuyển logic này sang C", hoặc một thư viện khác có thể xử lý các vòng lặp lồng nhau chỉ liên quan đến mảng không? Tôi tìm kiếm nó cho trường hợp chung, không phải giải pháp thay thế cho ví dụ cụ thể này (nhưng nếu bạn có, tôi có thể mở một phần Hỏi & Đáp riêng).

Answer 1

Về cơ bản đây là ý tưởng đằng sau Numba . Không nhanh như C, nhưng nó có thể gần ... Nó sử dụng trình biên dịch jit để biên dịch mã python sang máy và nó tương thích với hầu hết các chức năng của Numpy. (Trong tài liệu, bạn tìm thấy tất cả các chi tiết)

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

Trăn nguyên chất:

%%timeit

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

338 ms ± 12,6 ms mỗi vòng (trung bình ± std. Dev. Của 7 lần chạy, mỗi vòng 1 lần)

Với Numba:

%%timeit

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

12 ms ± 534 µs mỗi vòng (trung bình ± std. Dev. Của 7 lần chạy, mỗi lần 100 vòng)