Python - 具有多个 for 循环的多处理

Python - Multiprocessing with multiple for-loops

我知道还有其他关于这个主题的问题,所以很抱歉我不得不再问一次,但我对这个主题还很陌生,所以我无法让它工作。

我有四个 for 循环(嵌套),其中进行了某些代数计算(例如矩阵运算)。这些计算需要太多时间才能完成,所以我希望我可以使用多处理来加快速度。

代码如下。我在这里模拟了范围和矩阵大小,但在我的代码中这些范围确实被使用了(所以花费这么长时间并不奇怪)。直接复制粘贴代码应该就可以运行了。

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, :])

解决这个问题的最佳方法是什么?如果多处理是解决方案,如何为我的案例实现它(因为我无法弄清楚)。

提前致谢。

一种方法是将两个内部循环移动到一个以 aibi 作为参数并返回索引和结果的函数中。然后使用 multiprocessing.Pool.imap_unordered() 到 运行 ai, bi 对上的函数。像这样的东西(未经测试):

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()

或者将内部三个循环放入函数中,一次为 A_in 生成一个“行”。对它进行两种方式的分析,看看哪个更快。

诀窍在于设置列表(a_list、b_list 等)和 Y 矩阵。这取决于它们的特性(常数,quickly/slowly 计算,large/small,等等)。