Tại sao einsum của numpy lại chậm hơn các chức năng tích hợp của numpy?

Question

Tôi thường nhận được hiệu suất tốt từ chức năng einsum của numpy (và tôi thích cú pháp của nó). @ Ophion của câu trả lời cho this question cho thấy rằng - đối với các trường hợp được kiểm tra - einsum luôn tốt hơn các chức năng "tích hợp sẵn" (đôi khi một chút, đôi khi rất nhiều). Nhưng tôi chỉ gặp phải một trường hợp mà einsum chậm hơn nhiều. Hãy xem xét các chức năng tương đương như sau:

(M, K) = (1000000, 20) 
C = np.random.rand(K, K) 
X = np.random.rand(M, K) 

def func_dot(C, X): 
    Y = X.dot(C) 
    return np.sum(Y * X, axis=1) 

def func_einsum(C, X): 
    return np.einsum('ik,km,im->i', X, C, X) 

def func_einsum2(C, X): 
    # Like func_einsum but break it into two steps. 
    A = np.einsum('ik,km', X, C) 
    return np.einsum('ik,ik->i', A, X) 

tôi mong đợi func_einsum để chạy nhanh nhất nhưng đó không phải là những gì tôi gặp phải. Chạy trên một CPU quad-core với hyperthreading, numpy phiên bản 1.9.0.dev-7ae0206, và đa luồng với OpenBLAS, tôi nhận được kết quả như sau:

In [2]: %time y1 = func_dot(C, X) 
CPU times: user 320 ms, sys: 312 ms, total: 632 ms 
Wall time: 209 ms 
In [3]: %time y2 = func_einsum(C, X) 
CPU times: user 844 ms, sys: 0 ns, total: 844 ms 
Wall time: 842 ms 
In [4]: %time y3 = func_einsum2(C, X) 
CPU times: user 292 ms, sys: 44 ms, total: 336 ms 
Wall time: 334 ms 

Khi tôi tăng K-200, sự khác biệt là cực đoan hơn :

In [2]: %time y1= func_dot(C, X) 
CPU times: user 4.5 s, sys: 1.02 s, total: 5.52 s 
Wall time: 2.3 s 
In [3]: %time y2= func_einsum(C, X) 
CPU times: user 1min 16s, sys: 44 ms, total: 1min 16s 
Wall time: 1min 16s 
In [4]: %time y3 = func_einsum2(C, X) 
CPU times: user 15.3 s, sys: 312 ms, total: 15.6 s 
Wall time: 15.6 s 

Ai đó có thể giải thích tại sao einsum chậm hơn nhiều ở đây?

Nếu vấn đề, đây là cấu hình NumPy tôi:

In [6]: np.show_config() 
lapack_info: 
    libraries = ['openblas'] 
    library_dirs = ['/usr/local/lib'] 
    language = f77 
atlas_threads_info: 
    libraries = ['openblas'] 
    library_dirs = ['/usr/local/lib'] 
    define_macros = [('ATLAS_WITHOUT_LAPACK', None)] 
    language = c 
    include_dirs = ['/usr/local/include'] 
blas_opt_info: 
    libraries = ['openblas'] 
    library_dirs = ['/usr/local/lib'] 
    define_macros = [('ATLAS_INFO', '"\\"None\\""')] 
    language = c 
    include_dirs = ['/usr/local/include'] 
atlas_blas_threads_info: 
    libraries = ['openblas'] 
    library_dirs = ['/usr/local/lib'] 
    define_macros = [('ATLAS_INFO', '"\\"None\\""')] 
    language = c 
    include_dirs = ['/usr/local/include'] 
lapack_opt_info: 
    libraries = ['openblas', 'openblas'] 
    library_dirs = ['/usr/local/lib'] 
    define_macros = [('ATLAS_WITHOUT_LAPACK', None)] 
    language = f77 
    include_dirs = ['/usr/local/include'] 
lapack_mkl_info: 
    NOT AVAILABLE 
blas_mkl_info: 
    NOT AVAILABLE 
mkl_info: 
    NOT AVAILABLE 

Answer 1

Bạn có thể có tốt nhất của cả hai thế giới:

def func_dot_einsum(C, X): 
    Y = X.dot(C) 
    return np.einsum('ij,ij->i', Y, X) 

Trên hệ thống của tôi:

In [7]: %timeit func_dot(C, X) 
10 loops, best of 3: 31.1 ms per loop 

In [8]: %timeit func_einsum(C, X) 
10 loops, best of 3: 105 ms per loop 

In [9]: %timeit func_einsum2(C, X) 
10 loops, best of 3: 43.5 ms per loop 

In [10]: %timeit func_dot_einsum(C, X) 
10 loops, best of 3: 21 ms per loop 

Khi có sẵn , np.dot sử dụng BLAS, MKL hoặc bất kỳ thư viện nào bạn có. Vì vậy, các cuộc gọi đến np.dot là gần như chắc chắn là đa luồng. np.einsum có các vòng lặp riêng của nó, do đó không sử dụng bất kỳ tối ưu nào, ngoài việc sử dụng SIMD của chính nó để tăng tốc độ thực hiện trên thực hiện vanilla C.

---

Sau đó, cuộc gọi einsum đa đầu vào chạy chậm hơn nhiều ... Nguồn yếu cho einsum rất phức tạp và tôi không hiểu đầy đủ. Vì vậy, được thông báo rằng sau đây là đầu cơ tốt nhất, nhưng đây là những gì tôi nghĩ là đang xảy ra ...

Khi bạn chạy một cái gì đó như np.einsum('ij,ij->i', a, b), lợi ích hơn làm np.sum(a*b, axis=1) xuất phát từ việc tránh phải khởi tạo mảng trung gian với tất cả sản phẩm và lặp lại hai lần trên đó. Vì vậy, ở cấp thấp những gì diễn ra trên một cái gì đó như:

for i in range(I): 
    out[i] = 0 
    for j in range(J): 
     out[i] += a[i, j] * b[i, j] 

Say bây giờ mà bạn đang theo đuổi một cái gì đó như:

np.einsum('ij,jk,ik->i', a, b, c) 

Bạn có thể làm các hoạt động tương tự như

np.sum(a[:, :, None] * b[None, :, :] * c[:, None, :], axis=(1, 2)) 

Và những gì tôi nghĩ rằng einsum thực hiện là chạy mã cuối cùng này mà không cần phải khởi tạo mảng trung gian lớn, điều này chắc chắn làm mọi thứ nhanh hơn:

In [29]: a, b, c = np.random.rand(3, 100, 100) 

In [30]: %timeit np.einsum('ij,jk,ik->i', a, b, c) 
100 loops, best of 3: 2.41 ms per loop 

In [31]: %timeit np.sum(a[:, :, None] * b[None, :, :] * c[:, None, :], axis=(1, 2)) 
100 loops, best of 3: 12.3 ms per loop 

Nhưng nếu bạn nhìn kỹ, việc loại bỏ lưu trữ trung gian có thể là một điều khủng khiếp.Đây là những gì tôi nghĩ einsum đang làm ở mức thấp:

for i in range(I): 
    out[i] = 0 
    for j in range(J): 
     for k in range(K): 
      out[i] += a[i, j] * b[j, k] * c[i, k] 

Nhưng bạn đang lặp lại một tấn hoạt động! Nếu bạn thay vì đã làm:

for i in range(I): 
    out[i] = 0 
    for j in range(J): 
     temp = 0 
     for k in range(K): 
      temp += b[j, k] * c[i, k] 
     out[i] += a[i, j] * temp 

bạn sẽ làm I * J * (K-1) phép nhân ít (và thêm I * J bổ sung), và tiết kiệm cho mình một tấn thời gian. Tôi đoán là einsum không đủ thông minh để tối ưu hóa mọi thứ ở cấp độ này. Trong các source code có một gợi ý rằng nó chỉ tối ưu hóa hoạt động với 1 hoặc 2 toán hạng, không 3. Trong mọi trường hợp tự động hóa điều này cho đầu vào chung có vẻ như bất cứ điều gì nhưng đơn giản ...

Answer 2

einsum có trường hợp đặc biệt cho '2 toán hạng , ndim = 2 '. Trong trường hợp này có 3 toán hạng, và tổng cộng 3 chiều. Vì vậy, nó phải sử dụng chung nditer.

Trong khi cố gắng để hiểu làm thế nào đầu vào chuỗi được phân tách, tôi đã viết một mô phỏng Python einsum tinh khiết, https://github.com/hpaulj/numpy-einsum/blob/master/einsum_py.py

The (rút gọn) einsum và phụ phẩm tổng của chức năng là:

def myeinsum(subscripts, *ops, **kwargs): 
    # dropin preplacement for np.einsum (more or less) 
    <parse subscript strings> 
    <prepare op_axes> 
    x = sum_of_prod(ops, op_axes, **kwargs) 
    return x 

def sum_of_prod(ops, op_axes,...): 
    ... 
    it = np.nditer(ops, flags, op_flags, op_axes) 
    it.operands[nop][...] = 0 
    it.reset() 
    for (x,y,z,w) in it: 
     w[...] += x*y*z 
    return it.operands[nop] 

gỡ rối đầu ra cho myeinsum('ik,km,im->i',X,C,X,debug=True) với (M,K)=(10,5)

{'max_label': 109, 
'min_label': 105, 
'nop': 3, 
'shapes': [(10, 5), (5, 5), (10, 5)], 
....}} 
... 
iter labels: [105, 107, 109],'ikm' 

op_axes [[0, 1, -1], [-1, 0, 1], [0, -1, 1], [0, -1, -1]] 

Nếu bạn viết một hàm sum-of-prod như tôi này n cython bạn sẽ nhận được một cái gì đó gần với tổng quát einsum.

Với đầy đủ (M,K), mô phỏng einsum này chậm hơn 6-7x.

---

Một số timings xây dựng trên câu trả lời khác:

In [84]: timeit np.dot(X,C) 
1 loops, best of 3: 781 ms per loop 

In [85]: timeit np.einsum('ik,km->im',X,C) 
1 loops, best of 3: 1.28 s per loop 

In [86]: timeit np.einsum('im,im->i',A,X) 
10 loops, best of 3: 163 ms per loop 

này 'im,im->i' step is substantially faster than the other. The sum dimension, m is only 20. I suspect einsum` được điều trị này như là một trường hợp đặc biệt.

In [87]: timeit np.einsum('im,im->i',np.dot(X,C),X) 
1 loops, best of 3: 950 ms per loop 

In [88]: timeit np.einsum('im,im->i',np.einsum('ik,km->im',X,C),X) 
1 loops, best of 3: 1.45 s per loop 

Thời gian cho các phép tính tổng hợp này chỉ đơn giản là tổng của các phần tương ứng.