使用多处理嵌套循环

Question

我有一个关于python中多处理的快速问题。

我正在对三个参数进行相当大的网格搜索，计算需要大约14个小时才能完成。 我想通过使用多处理来减少运行时间。

我的代码的一个非常简化的示例在这里：

import numpy as np
import pickle
import time

a_range = np.arange(14, 18, 0.2)
b_range = np.arange(1000, 5000, 200)
c_range = np.arange(12, 21, .5)

a_position = range(len(a_range))
b_position = range(len(b_range))
c_position = range(len(c_range))

data_grid = np.zeros([len(a_range), len(b_range), len(c_range)])
record_data = []

start_time = time.time()

for (a,apos) in zip(a_range, a_position):
    for (b, bpos) in zip(b_range, b_position):
        for (c, cpos) in zip(c_range, c_position):
            example = a+b+c  #The math in my model is much more complex and takes
            #about 7-8 seconds to process
            data_grid[apos, bpos, cpos] = example
            record_data.append([a, b, c, example])

with open('Test_File', 'wb') as f: 
    pickle.dump(record_data, f) 

np.save('example_values', data_grid) 

print 'Code ran for ', round(time.time()-start_time,2), ' seconds'

现在，我在多处理方面的经验绝对为零，因此我的第一个尝试是将for循环更改为一个函数，然后像这样调用多处理函数：

def run_model(a, b, c, apos, bpos, cpos):
    example=a+b+c  
    data_grid[apos, bpos, cpos]=example
    record_data.append([a, b, c, example])

from multiprocessing import Pool

if __name__=='__main__':
    pool=Pool(processes=4)
    pool.map(run_model, [a_range, b_range, c_range, a_position, b_positon, c_positon])
    pool.close()
    pool.join()

但是，这在pool.map调用中失败了。 我知道此函数仅接受一个可迭代的参数，但不知道如何解决该问题。 我也对data_grid变量是否正确填充表示怀疑。 我要从此函数获得的结果是保存了两个文件，一个保存为值数组，其索引对应于a，b和c值，最后一个包含a，b，c值和结果值的列表列表（示例在上面的代码中）

谢谢你的帮助！

-将

Answer 1

这不能解决您的多处理问题，但是可以使您的处理更快。

您使用嵌套的循环来构建第二坐标，然后对他们的经营模式可以使用矢量化 d ```numpy.meshgrid```` 。 在不知道您的实际计算的情况下，这种方法无法进行测试。

import numpy as np
a = np.array([0,1,2])
b = np.array([10,11,12])
c = np.array([20,21,22])

x, y, z = np.meshgrid(a,b,c)

>>> x
array([[[0, 0, 0],
        [1, 1, 1],
        [2, 2, 2]],

       [[0, 0, 0],
        [1, 1, 1],
        [2, 2, 2]],

       [[0, 0, 0],
        [1, 1, 1],
        [2, 2, 2]]])
>>> y
array([[[10, 10, 10],
        [10, 10, 10],
        [10, 10, 10]],

       [[11, 11, 11],
        [11, 11, 11],
        [11, 11, 11]],

       [[12, 12, 12],
        [12, 12, 12],
        [12, 12, 12]]])
>>> z
array([[[20, 21, 22],
        [20, 21, 22],
        [20, 21, 22]],

       [[20, 21, 22],
        [20, 21, 22],
        [20, 21, 22]],

       [[20, 21, 22],
        [20, 21, 22],
        [20, 21, 22]]])
>>> 



f = x + y + z

>>> f
array([[[30, 31, 32],
        [31, 32, 33],
        [32, 33, 34]],

       [[31, 32, 33],
        [32, 33, 34],
        [33, 34, 35]],

       [[32, 33, 34],
        [33, 34, 35],
        [34, 35, 36]]])
>>> 

---

也可以选择使用meshgrid创建实际点，然后使用一个循环遍历这些点-除非您能弄清楚如何重塑结果，否则使用这种方法会丢失空间信息。 我在答案中找到了这个https://stackoverflow.com/a/18253506/2823755

points = np.vstack([x,y,z]).reshape(3, -1).T

>>> points
array([[ 0, 10, 20],
       [ 0, 10, 21],
       [ 0, 10, 22],
       [ 1, 10, 20],
       [ 1, 10, 21],
       [ 1, 10, 22],
       [ 2, 10, 20],
       [ 2, 10, 21],
       [ 2, 10, 22],
       [ 0, 11, 20],
       [ 0, 11, 21],
       [ 0, 11, 22],
       [ 1, 11, 20],
       [ 1, 11, 21],
       [ 1, 11, 22],
       [ 2, 11, 20],
       [ 2, 11, 21],
       [ 2, 11, 22],
       [ 0, 12, 20],
       [ 0, 12, 21],
       [ 0, 12, 22],
       [ 1, 12, 20],
       [ 1, 12, 21],
       [ 1, 12, 22],
       [ 2, 12, 20],
       [ 2, 12, 21],
       [ 2, 12, 22]])
>>>

您可以创建一个函数并将其应用于points

def g(point):
    x, y, z = point
    return x + y + z

result = np.apply_along_axis(g, 1, points)

>>> result
array([30, 31, 32, 31, 32, 33, 32, 33, 34, 31, 32, 33, 32, 33, 34, 33, 34, 35, 32, 33, 34, 33, 34, 35, 34, 35, 36])
>>>

重塑此示例非常简单：

>>> result.reshape(3,3,3)
array([[[30, 31, 32],
        [31, 32, 33],
        [32, 33, 34]],

       [[31, 32, 33],
        [32, 33, 34],
        [33, 34, 35]],

       [[32, 33, 34],
        [33, 34, 35],
        [34, 35, 36]]])
>>> 

测试以确保它们都相同

>>> np.all(result.reshape(3,3,3) == f)
True
>>> 

对于更复杂的数学，只需遍历点即可：

result = []
for point in points:
    example = some_maths
    result.append(example)

result = np.array(result).reshape(shape_of_the_3d_data)

Answer 2

根据每个用户wwii的建议，我通过使用numpy的meshgrid重写了上面的示例，并摆脱了仅用于单个循环的嵌套for循环。 这是工作代码的示例。

import numpy as np
import time

a_range = np.arange(14, 18, 1)
b_range = np.arange(1000, 2200, 200)
c_range = np.arange(12, 21, 1)

a_position = range(len(a_range))
b_position = range(len(b_range))
c_position = range(len(c_range))

mesha, meshb, meshc = np.meshgrid(a_range, b_range, c_range)
mesh_vals = np.vstack([mesha, meshb, meshc]).reshape(3, -1).T

mesha_pos, meshb_pos, meshc_pos = np.meshgrid(a_position, b_position, c_position)
mesh_positions = np.vstack([mesha_pos, meshb_pos, meshc_pos]).reshape(3,-1).T

data_grid = np.zeros([len(a_range), len(b_range), len(c_range)])
record_data = []

start_time = time.time()

for pol in range(len(mesh_positions)):
    example = mesh_vals[pol][0]+ mesh_vals[pol][1]+ mesh_vals[pol][2]
    data_grid[mesh_positions[pol][0], mesh_positions[pol][1], mesh_positions[pol][2]] = example
    record_data.append([mesh_vals[pol][0], mesh_vals[pol][1], mesh_vals[pol][2], example])

print 'Code ran for ', round(time.time()-start_time,2), ' seconds'

实际上，在进一步调查之后，这导致运行时间显着增加。 对于a，b和c而言，for循环与该方法之间的差异为20秒。 我不知道为什么，但是我确实知道问题的这种构造应该使多处理更容易，因为只有一个for循环可以处理。