Python threading parallelizing simple operations over 2D grid array

Viewed 56

I have the following code structure which I want to take advantage of multi-threading to speed up results, as every computation is independent. For every point in a 2D grid, I call a compute_at_point() function and get the result.

How can I parallelize this with threads? I know the computers that will be used have at least 4 cores. Code below:

for i in range(0, grid_rows):
        for j in range(0, grid_cols):

            grid_point = input_grid[i, j]
            res_at_point = compute_at_point(input_grid, grid_point)
            output_grid[i, j] = res_at_point

Where input_grid and output_grid have the same shape, and every computation is independent.

1 Answers

Using multiprocessing.Pool and functools.partial, we can do the following:

import numpy as np
from functools import partial
from multiprocessing import Pool

input_grid = np.random.random((3,4))
grid_rows = input_grid.shape[0]
grid_cols = input_grid.shape[1]

def compute_at_point(input_grid, grid_point):
    return np.mean(input_grid) + grid_point

with Pool(processes=4) as p:
    output_grid = p.map(partial(compute_at_point, input_grid), input_grid.flatten())
    output_grid = np.array(output_grid).reshape(grid_rows, grid_cols)

The simplest way to solve multiprocessing problems is by having one iterable, and allowing the the multiprocessing library to iterate over it. In this instance, we could pass the input_grid indices, or just flatten the array and iterate over its values.

Related