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) 

我将索引结果，以尝试使用以下方法恢复所有生成的cube：

results[:,0].shape

在这个例子中，我得到的结果：

(10,)

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

编辑：在这个例子中，有必要返回a，b和c，这是一个较大的代码段的简单示例，我需要返回一个立方体和多个参数。

提前非常感谢！

Answer 1

我期望在以下示例中(10,5,3,3)

要获得最终数组的形状，您不需要将这些变量拖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)

Answer 2

当我运行这段代码时，我得到一个返回的(10,4)numpy 数组，它对应于(ntrials, len(outputTuple)。因此，return cube, a, b, c被打包为元组，然后将其放置在results数组中。因此，如果您想访问您的值：

cube = results[:,0] # returns a (10, 5, 3, 3) as you want
a = results[:,1] # (10,) array with the a values
b = results[:,2] # (10,) array
c = results[:,3] # (10,) array

或者，如果您希望将所有辅助参数放在一起：

abc = results[:,1:] # (10, 3) array. Aux parameters separated by trial