Python 多进程 - 索引多个返回

Question

我正在尝试使用 Python 的多处理池功能返回多维数组以及一些元参数。

但是，当我尝试索引多维数组并检查其大小时，我得到的大小为 (N,)，而不是我期望的大小，在以下示例中为 (10,5,3,3):

import multiprocessing as mp
import numpy as np
from tqdm import tqdm

def function(x):
        cube = np.ones((5,3,3))
        a,b,c = 1,2,3
        return cube,a,b,c

pool = mp.Pool(processes=4)

results = list(tqdm(pool.imap(function,range(10)),total=10))
results = [x for x in results if str(x) != 'nan']
results = np.array(results) 

我将使用以下命令对结果进行索引，以尝试恢复所有生成的多维数据集:

results[:,0].shape

在这个例子中我得到了结果:

(10,)

我觉得这是一个相当基本的问题，但是有没有办法以索引结果产生我期望看到的多维形状的方式设置这个多处理代码？

编辑:在此示例中必须返回 a、b 和 c，这是一个较大代码段的简单示例，我需要返回一个多维数据集和多个参数。

提前非常感谢!

Answer 1

what I expect which is (10,5,3,3) in the following example

要获得最终数组的形状，您不需要拖动这些变量a,b,c作为目标函数的结果，只需返回cube (多维 numpy 数组)。实际上，它们在这种情况下似乎没有意义。

import multiprocessing as mp
import numpy as np
from tqdm import tqdm

def function(x):
        cube = np.ones((5,3,3))
        # a,b,c = 1,2,3
        return cube

pool = mp.Pool(processes=4)

results = list(tqdm(pool.imap(function,range(10)),total=10))
results = [x for x in results if str(x) != 'nan']
results = np.array(results)
print(results.shape)

输出:

100%|██████████| 10/10 [00:00<00:00, 15845.50it/s]
(10, 5, 3, 3)

<小时/>

如果需要返回多个变量 - 只需从结果中提取所有立方体:

import multiprocessing as mp
import numpy as np
from tqdm import tqdm

def function(x):
        cube = np.ones((5,3,3))
        a,b,c = 1,2,3
        return cube, a, b, c

pool = mp.Pool(processes=4)

results = list(tqdm(pool.imap(function,range(10)),total=10))
results = [x for x in results if str(x) != 'nan']
cubes = np.array([r[0] for r in results])
print(results[0])  # print 1st result item
print(cubes.shape)

输出:

100%|██████████| 10/10 [00:00<00:00, 51590.46it/s]
(array([[[1., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.]],

       [[1., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.]],

       [[1., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.]],

       [[1., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.]],

       [[1., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.]]]), 1, 2, 3)
(10, 5, 3, 3)