并行化数百万次 Numpy 函数的迭代

Question

I'm running the below function compare on 20 million different combinations of the parameter sample where sample is a 1D array composed of 100 1s and 0s.

compare takes a couple of other arrays along with sample and uses them to perform a few dot products, exponentiate those dot products, and then compare them relative to each other. These other arrays stay the same.

On my laptop, it takes about an hour to run through all 20 million combinations.

I'm looking for ways to make it go quicker. I am open to both improving the below code and using libraries such as dask which take advantage of parallel processing.

Notes:

• The comments on each line in compare show a very rough estimate of how long that line takes on my machine. They're the result of a %%timeit on the line on its own outside the function.
• The inputs to compare are not actually randomly generated in my use case

def compare(sample, competition_exp_dot, preferences): # 140 µs
    sample_exp_dot = np.exp(preferences @ sample) #30.3 µs
    all_competitors = np.append(sample_exp_dot.reshape(-1, 1), competition_exp_dot, 1) # 5 µs
    all_results = all_products/all_competitors.sum(axis=1)[:,None] #27.4 µs

    return all_results.mean(axis=0) #20.6 µs

#these inputs to the above function stay the same
preferences = np.random.random((1000,100))
competition = np.array([np.random.randint(0,2,100), np.random.randint(0,2,100)])
competition_exp_dot = np.exp(preferences @ competition.T)

# the function is run with 20,000,000 variations of sample
population = np.random.randint(0,2,(20000000,100))
result = [share_calc(sample, competition_exp_dot, preferences) for sample in population]

Answer 1

There are many ways to accelerate simple array programming code like this:

1. You can use a tool like Numba, which will fuse some of the work, and also provide some options for single-node multi-core parallelism
2. You can use a tool like Dask to scale this onto multiple cores of a single machine (also possible with Numba) or across a cluster
3. You can use one of the GPU array libraries, like Torch, TensorFlow, CuPy, or Jax to run this on a GPU

You can also do any mixture of the above.

Answer 2

You can consider the following:

import torch
import numpy as np
x = np.array([[1,2,3],[4,5,6]])
b = torch.from_numpy(x)
if torch.cuda.is_available():
    device = torch.device("cuda")
b = b.to(device)

Answer 3

I implemented numba as suggested by MRocklin. The result is about 4 times faster on my machine.

Modified Numba Version

import numba as nb

@nb.jit
def nb_compare(sample, competition_exp_dot, preferences):
    sample_exp_dot = np.exp(preferences @ sample)
    all_competitors = np.append(sample_exp_dot.reshape(-1, 1), competition_exp_dot, 1)
    all_results = (all_competitors.T/all_competitors.sum(axis=1)).T

    return np_mean(all_results, 0) # see source for np_mean in notes below

Comparable Numpy Version

import numpy as np
def np_compare(sample, competition_exp_dot, preferences):
    sample_exp_dot = np.exp(preferences @ sample)
    all_competitors = np.append(sample_exp_dot.reshape(-1, 1), competition_exp_dot, 1)
    all_results = (all_competitors.T/all_competitors.sum(axis=1)).T

    return all_results.mean(axis=0)

Timing Comparison

Setup:

preferences = np.random.random((1000,100)).astype(np.float32)
competition = np.array([np.random.randint(0,2,100), np.random.randint(0,2,100)]).astype(np.float32)
competition_exp_dot = np.exp(preferences @ competition.T)
sample = np.random.randint(0,2,100)

%timeit np_compare(sample, competition_exp_dot, preferences)
"210 µs ± 13.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)"


%timeit -n 10000 nb_compare(population[0], competition_exp_dot, preferences)
"52.4 µs ± 4.48 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)"

Notes

Numba doesn't support optional parameters like axis for np.mean and returns a TypingError. In my numba code, I use call below version of np_mean instead.

Credit to joelrich

import numba as nb, numpy as np

# fix to use np.mean along axis=0 (numba doesn't support optional arguments for np.mean)
# credit to: joelrich https://github.com/numba/numba/issues/1269#issuecomment-472574352
@nb.njit
def np_apply_along_axis(func1d, axis, arr):
  assert arr.ndim == 2
  assert axis in [0, 1]
  if axis == 0:
    result = np.empty(arr.shape[1])
    for i in range(len(result)):
      result[i] = func1d(arr[:, i])
  else:
    result = np.empty(arr.shape[0])
    for i in range(len(result)):
      result[i] = func1d(arr[i, :])
  return result

@nb.njit
def np_mean(array, axis):
  return np_apply_along_axis(np.mean, axis, array)