python - 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的 s，使用以下命令：

results[:,0].shape

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

(10,)

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

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

提前谢谢了！

标签： pythonarraysindexingmultiprocessing

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)