Multiprocesamiento Python para operaciones costosas con matriz 2D

Tengo una función, digamos fun(a,b) que es bastante cara y devuelve un conjunto de datos.

Mi método actual se ve así:

 a = np.linspace(0,100,300) b = np.linspace(0,100,300) A,B = np.meshgrid(a,b) Y = np.zeros(A.shape) for i,av in enumerate(a): for j, bv in enumerate(b): Y[i,j] = fun(av,bv) 

(Algo así, tuve que confundir un poco las cosas para que encajaran). De todos modos, el proceso lleva bastante tiempo, y me preguntaba si había una forma sencilla de usar mi procesador multinúcleo para acelerar las cosas.

Hay un gran módulo llamado multiprocesamiento que forma parte de la biblioteca estándar de python. Generará procesos en todos los núcleos que desee para aprovechar las otras CPU. Hay un ejemplo del uso del objeto Pool en los documentos, a continuación se muestra una versión abreviada de ese ejemplo. Calculará el cuadrado de 10 números distribuyendo la carga de trabajo para los procesos de trabajo y mostrando el resultado.

Grupo de trabajadores simple

 from multiprocessing import Pool def f(x): return x*x pool = Pool(processes=4) print pool.map(f, range(10)) 

Salida

 [0, 1, 4, 9, 16, 25, 36, 49, 64, 81] 

Tuve un desafío mayor para dividir su problema en el mismo tipo de estructura. Tuve que crear algunas funciones intermedias para lograr el trabajo. No tengo números, así que solo usé listas y diccionarios en lugar de lo que pusiste. Puedes reemplazarlos y probar el código.

Escenario más complicado

 from multiprocessing import Pool import time, pprint def fun(av, bv): time.sleep(0.1) return (av, bv) def data_stream(a, b): for i, av in enumerate(a): for j, bv in enumerate(b): yield (i, j), (av, bv) def proxy(args): return args[0], fun(*args[1]) a = range(100, 400, 100) b = range(100, 400, 100) Y = {} pool = Pool(processes=4) results = pool.map(proxy, data_stream(a, b)) for k,v in results: Y[k] = v pprint.pprint(Y) 

salida

 {(0, 0): (100, 100), (0, 1): (100, 200), (0, 2): (100, 300), (1, 0): (200, 100), (1, 1): (200, 200), (1, 2): (200, 300), (2, 0): (300, 100), (2, 1): (300, 200), (2, 2): (300, 300)} 

Actuación

En el ejemplo, acabo de poner una demora de 0.1 segundos para simular un trabajo pesado. Pero incluso en este ejemplo, si ejecuta un grupo con processes=1 se ejecuta en 0.950s con processes=4 se ejecuta en 0.352s. Puede usar la biblioteca de multiprocesamiento de muchas maneras diferentes. Pool es solo de una manera. Es posible que desee explorar los ejemplos y experimentar.

En uno de los comentarios a continuación, se mencionó el uso del argumento chunksize para pool.map para ayudar a mejorar el rendimiento. Es importante tener una comprensión general de lo que está pasando debajo del capó para realmente controlar el rendimiento. Básicamente, todos los datos que está pasando a los otros procesos deben ser decapados, y luego el resultado pasa por el mismo proceso al proceso principal. Hay una sobrecarga a esta comunicación entre procesos. Ten esto en cuenta cuando hagas experimentos.