python 多处理池示例，使用3D numpy数组作为输入，在计算密集型堆栈上使用所有CPU内核。

Posted 2021-05-09

from multiprocessing import Pool
import numpy as np


n = 1000000
x = 5
y = 5
z = 100


def f(x):
    #print x.shape
    for i in range(n):
        x = x**2
        x = np.sqrt(x)


    return np.sum(x)


if __name__ == '__main__':
    m = np.random.random((x,y,z))
    #print m[:,:].shape
    
    result = np.ndarray((x,y), dtype=object)
    
    p = Pool()
    for i in range(m.shape[0]):
        for j in range(m.shape[1]):
            result[i,j] = p.apply_async(f, [m[i,j,:]])
    p.close()
    
    numresult = np.zeros((x,y))
    
    for i in range(m.shape[0]):
        for j in range(m.shape[1]):
            numresult[i,j] = result[i,j].get()
                    
    print numresult