Supongamos que tengo dos vectores y deseo obtener su producto escalar; esto es simple,
import numpy as np a = np.random.rand(3) b = np.random.rand(3) result = np.dot(a,b)Si tengo montones de vectores y quiero que cada uno tenga puntos, el código más ingenuo es
# 5 = number of vectors a = np.random.rand(5,3) b = np.random.rand(5,3) result = [np.dot(aa,bb) for aa, bb in zip(a,b)]Dos formas de procesar por lotes este cálculo son mediante multiplicar y sumar, y einsum,
result = np.sum(a*b, axis=1) # or result = np.einsum('ij,ij->i', a, b) Sin embargo, ninguno de estos se envía al backend de BLAS y, por lo tanto, usa solo un núcleo. Esto no es muy bueno cuando N es muy grande, digamos 1 millón.
tensordot se envía al backend de BLAS. Una forma terrible de hacer este cálculo con tensordot es
np.diag(np.tensordot(a,b, axes=[1,1]) Esto es terrible porque asigna una matriz N*N , y la mayoría de los elementos son trabajo inútil.
Otro enfoque (brillantemente rápido) es la función oculta inner1d
from numpy.core.umath_tests import inner1d result = inner1d(a,b)pero parece que esto no va a ser viable , ya que el tema que podría exportarlo públicamente se ha vuelto obsoleto. Y esto todavía se reduce a escribir el ciclo en C, en lugar de usar varios núcleos.
¿Hay alguna manera de hacer que dot , matmul o tensordot hagan todos estos productos de punto a la vez, en múltiples núcleos?
En primer lugar, no hay una función BLAS directa para hacer eso . El uso de muchas llamadas a la función BLAS de nivel 1 no es muy eficiente, ya que el uso de varios subprocesos para un cálculo de tiempo muy corto tiende a introducir una sobrecarga bastante grande y no usar varios subprocesos puede ser subóptimo. Aún así, dicho cálculo está principalmente ligado a la memoria y, por lo tanto, se escala mal en plataformas con muchos núcleos (pocos núcleos suelen ser suficientes para saturar el ancho de banda de la memoria).
Una solución simple es usar el paquete Numexpr que debería hacerlo de manera bastante eficiente (debería evitar la creación de matrices temporales y también debería usar múltiples subprocesos). Sin embargo, el rendimiento es un tanto decepcionante para arreglos grandes en este caso.
La mejor solución parece usar Numba (o Cython). Numba puede generar un código rápido para matrices de entrada grandes y pequeñas y es fácil paralelizar el código. Sin embargo, tenga en cuenta que la administración de subprocesos presenta una sobrecarga que puede ser bastante grande para arreglos pequeños (hasta unos pocos ms en algunas plataformas de muchos núcleos).
Aquí hay una implementación de Numexpr:
import numexpr as ne expr = ne.NumExpr('sum(a * b, axis=1)') result = expr.run(a, b)Aquí hay una implementación (secuencial) de Numba:
import numba as nb # Use `parallel=True` for a parallel implementation @nb.njit('float64[:](float64[:,::1], float64[:,::1])') def multiDots(a, b): assert a.shape == b.shape n, m = a.shape res = np.empty(n, dtype=np.float64) # Use `nb.prange` instead of `range` to run the loop in parallel for i in range(n): s = 0.0 for j in range(m): s += a[i,j] * b[i,j] res[i] = s return res result = multiDots(a, b)Aquí hay algunos puntos de referencia en una (antigua) máquina de 2 núcleos:
On small 5x3 arrays: np.einsum('ij,ij->i', a, b, optimize=True): 45.2 us Numba (parallel): 12.1 us np.sum(a*b, axis=1): 9.5 us np.einsum('ij,ij->i', a, b): 6.5 us Numexpr: 3.2 us inner1d(a, b): 1.8 us Numba (sequential): 1.3 us On small 1000000x3 arrays: np.sum(a*b, axis=1): 27.8 ms Numexpr: 15.3 ms np.einsum('ij,ij->i', a, b, optimize=True): 9.0 ms np.einsum('ij,ij->i', a, b): 8.8 ms Numba (sequential): 6.8 ms inner1d(a, b): 6.5 ms Numba (parallel): 5.3 ms La implementación secuencial de Numba ofrece una buena compensación. Puede usar un interruptor si realmente desea obtener el mejor rendimiento. Sin embargo, elegir el mejor umbral n de forma independiente de la plataforma no es tan fácil.