Multiprocessing a function and adding results to a commonly shared array

user_11022471 0
Viewed 77

I have a dummy function test_func(n), which multiplies the np.ones (i.e., a 3x3x3 array) with a number n.

def test_func(n):
    array = n*np.ones((3,3,3))
    return array

Let's say I want to run this function in parallel and want the total sum of all the arrays, while keeping the array's shape intact (i.e., the output array should also be 3x3x3). For this, I normally define a variable that can save the results of the multiprocessing simulations and then add them all up (final_array).

#Multiprocessing starts here
from multiprocessing import Pool
if __name__ == '__main__':
    pool = Pool()
    grid = np.arange(2,4,1) # generates array [2,3]
    final_array  = pool.map(test_func, grid) # Holds the results of the multiprocessing simulations
    final_array  = sum(final_array, axis = 0) # A bit indirect
    pool.close()
    pool.join()

print(final_array) 
#[Output]: This is what I want from the code in the next column
array([[[5., 5., 5.],
        [5., 5., 5.],
        [5., 5., 5.]],

       [[5., 5., 5.],
        [5., 5., 5.],
        [5., 5., 5.]],

       [[5., 5., 5.],
        [5., 5., 5.],
        [5., 5., 5.]]])

Although a bit indirect (code shown above), it gets the job done. But a not-so-recent discussion on StackOverflow showed me that this could be done in a more elegant way: https://stackoverflow.com/questions/7894791/use-numpy-array-in-shared-memory-for-multiprocessing (answered by @jfs).

I tried to replicate the process but am unsure what I am doing wrong (trial code below).

import ctypes
import numpy as np

L, N, M = 3,3,3 # Shape of the array
mp_arr = mp.Array(ctypes.c_double, L * N * M)
final_array = np.frombuffer(mp_arr.get_obj())
final_array = final_array.reshape((L, N, M))

def test_func(n):
    final_array = np.frombuffer(mp_arr.get_obj())
    final_array = TC_p_value.reshape((L, N, M))
    final_array += n*np.ones((3,3,3))

def init(shared_arr_): # I do not even know what his one does
    global mp_arr
    mp_arr = shared_arr_

#Multiprocessing starts here
from multiprocessing import Pool
if __name__ == '__main__':
    pool = Pool(initializer=init, initargs=(mp_arr,))
    grid = np.arange(2,4,1)
    pool.map_async(test_func, grid)
    pool.close()
    pool.join()
shared-memory numpy function multiprocessing python
1 Answers

First, initialize the shared Array inside __main__ block. Also, use mp_arr.get_lock() because every process is accessing the whole 3x3x3 array:

import ctypes
import numpy as np
from multiprocessing import Pool, Array


def test_func(n):
    with mp_arr.get_lock():  # <-- every process is accessing the same slice, so get_lock()
        final_array = np.frombuffer(mp_arr.get_obj())
        final_array = final_array.reshape((L, N, M))
        final_array += n * np.ones((3, 3, 3))


def init(shared_arr_):
    global mp_arr
    mp_arr = shared_arr_


if __name__ == "__main__":
    L, N, M = 3, 3, 3
    mp_arr = Array(ctypes.c_double, L * N * M)
    grid = np.arange(2, 4, 1)

    with Pool(initializer=init, initargs=(mp_arr,)) as pool:
        pool.map(test_func, grid)

    final_array = np.frombuffer(mp_arr.get_obj())
    final_array = final_array.reshape((L, N, M))
    print(final_array)

Prints:

[[[5. 5. 5.]
  [5. 5. 5.]
  [5. 5. 5.]]

 [[5. 5. 5.]
  [5. 5. 5.]
  [5. 5. 5.]]

 [[5. 5. 5.]
  [5. 5. 5.]
  [5. 5. 5.]]]
user_100359850
Related