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 矩阵。 这取决于它们的特性（恒定、快速/缓慢计算、大/小等）。