Python - 具有多個 for 循環的多處理

Question

我知道還有其他關於這個話題的問題，所以很抱歉我不得不再問一次，但我無法讓它發揮作用，因為我對這個話題很陌生。

我有四個 for 循環（嵌套），其中完成了某些代數計算（例如矩陣運算）。 這些計算需要太多時間才能完成，所以我希望我可以通過 Multiprocessing 加快計算速度。

代碼如下。 我在這里模擬了范圍和矩陣大小，但在我的代碼中，這些范圍確實被使用了（所以花這么長時間並不奇怪）。 您應該能夠在復制粘貼代碼時直接運行它。

import numpy as np
from scipy.linalg import fractional_matrix_power
import math

#Lists for the loop (and one value)
x_list = np.arange(0, 32, 1)
y_list = np.arange(0, 32, 1)
a_list = np.arange(0, 501, 1)
b_list = np.arange(0, 501, 1)
c_list = np.arange(0, 64, 1)
d_number = 32

#Matrices
Y = np.arange(2048).reshape(32, 64)
g = np.asmatrix(np.empty([d_number, 1], dtype=np.complex_))
A = np.empty([len(a_list), len(b_list), len(c_list)], dtype=np.complex_)
A_in = np.empty([len(a_list), len(b_list)], dtype=np.complex_)


for ai in range(len(a_list)):
    
    for bi in range(len(b_list)):
        
        for ci in range(len(c_list)):
            
            f_k_i = c_list[ci]
            X_i = np.asmatrix([Y[:, ci]]).T
            
            for di in range(d_number):
                
                r = math.sqrt((x_list[di] - a_list[ai])**2 + (y_list[di] - b_list[bi])**2 + 63**2)
                g[di, 0] = np.exp(-2 * np.pi * 1j * f_k_i * (r / 8)) / r #g is a vector
            
            A[-bi, -ai, ci] = ((1 / np.linalg.norm(g)**2) * (((g.conj().T * fractional_matrix_power((X_i * X_i.conj().T), (1/5)) * g) / np.linalg.norm(g)**2)**2)).item(0)
         
        A_in[-bi, -ai] = (1 / len(c_list)) * sum(A[-bi, -ai, :])

解決這個問題的最佳方法是什么？ 如果多處理是解決方案，如何為我的案例實現這個（因為我無法弄清楚）。

提前致謝。

Answer 1

解決它的一種方法是將兩個內部循環移動到 function 中，將ai和bi作為參數並返回索引和結果。 然后使用multiprocessing.Pool.imap_unordered()在ai 、 bi對上運行 function。 像這樣的東西（未經測試）：

def partial_calc(index):
    """
    This function replaces the inner two loops to calculate the value of
    A_in[-bi, -ai]. index is a tuple (ai, bi).
    """

    ai, bi = index

    for ci in range(len(c_list)):
        
        f_k_i = c_list[ci]
        X_i = np.asmatrix([Y[:, ci]]).T
        
        for di in range(d_number):
            
            r = math.sqrt((x_list[di] - a_list[ai])**2 + (y_list[di] - b_list[bi])**2 + 63**2)
            g[di, 0] = np.exp(-2 * np.pi * 1j * f_k_i * (r / 8)) / r #g is a vector
        
        A[-bi, -ai, ci] = ((1 / np.linalg.norm(g)**2) * (((g.conj().T * fractional_matrix_power((X_i * X_i.conj().T), (1/5)) * g) / np.linalg.norm(g)**2)**2)).item(0)
        
    return ai, bi, (1 / len(c_list)) * sum(A[-bi, -ai, :])
    

def main():

    with multiprocessing.Pool(None) as p:
        # this replaces the outer two loops
        indices = itertools.product(range(len(a_list)), range(len(b_list)))

        partial_results = p.imap_unordered(partial_calc, indices)        
         
        for ai, bi, value in partial_results:
            A_in[-bi, -ai] = value

    #... do something with A_in ...

if __name__ == "__main__":
    main()

或者將內部的三個循環放入 function 並一次為A_in生成一個“行”。 以兩種方式對其進行分析，看看哪個更快。

訣竅是設置列表（a_list、b_list 等）和 Y 矩陣。 這取決於它們的特性（恆定、快速/緩慢計算、大/小等）。