使用多处理嵌套循环

Question

I have a quick question regarding multiprocessing in python.

I am conducting a rather large grid search over three parameters and the computation is taking ~14 hours to complete. I would like to shrink this run time down by using multiprocessing.

A very simplified example of my code is here:

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'

Now, I have absolutely zero experience in multiprocessing so my first attempt at this was changing the for loops into a function and then calling the multiprocessing function like this:

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()

This failed however at the pool.map call. I understand this function only takes a single iterable argument but I don't know how to fix the problem. I am also skeptical that the data_grid variable is going to be filled correctly. The result I want from this function is two files saved, one as an array of values whose indexes correspond to a, b, and c values and the last a list of lists containing the a, b, c values and the resulting value (example in the code above)

Thanks for any help!

-Will

Answer 1

This doesn't solve your multiprocessing problem but it might make your process faster.

Your pattern of using nested loops to construct nd coordinates and then operating on them can be vectorize d using ```numpy.meshgrid```` . Without knowing your actual calcs this approach can't be tested.

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]]])
>>> 

---

There is also the option of using meshgrid to create the actual points then use a single loop to iterate over the points - you lose the spatial info with this approach unless you can figure out how to reshape the result. I found this in SO answer 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]])
>>>

You can create a function and apply it to 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])
>>>

Reshaping this example is straightforward:

>>> 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]]])
>>> 

Test to make sure they both the same

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

For more complicated maths, just iterate over points:

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

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

Answer 2

As per user wwii's suggestions, I have rewrote the example above by using numpy's meshgrid and getting rid of the nested for loops for just a single loop. Here is an example of the working code.

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'

This actually, after further investigation caused the run time to increase rather significantly. The difference between the for loops and this method was 20 seconds when subjected to a large range for a, b, and c. I have no idea why, but I do know this construction of the problem should make multiprocessing easier since there is only a single for loop to deal with.