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

Question

我在2000万个不同的参数组合上运行下面的函数compare，其中sample是由100个1和0组成的一维数组

compare将其他数组与sample一起使用，并使用它们执行一些点积，对这些点积求幂，然后将它们相互比较。其他数组保持不变。你知道吗

在我的笔记本电脑上，运行2000万个组合大约需要一个小时。你知道吗

我在想办法让它走得更快。我对改进下面的代码和使用像dask这样利用并行处理的库持开放态度。你知道吗

备注：

• compare中每一行上的注释显示了该行在我的机器上花费的时间的一个非常粗略的估计。它们是%%timeit在函数外部自己联机的结果。你知道吗
• 在我的用例中，compare的输入实际上不是随机生成的

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

我按照洛克林先生的建议实现了numba。结果在我的机器上快了4倍。你知道吗

修改的Numba版本

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

可比Numpy版本

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)

时间比较

设置：

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)"

注意事项

Numba不支持可选参数，比如axis fornp.平均值并返回一个打字机错误。在我的numba代码中，我使用了np_mean的callbelow版本。你知道吗

记入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)

Answer 2

您可以考虑以下几点：

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

有许多方法可以加速这样的简单数组编程代码：

1. 您可以使用像Numba这样的工具，它将融合一些工作，还为单节点多核并行提供一些选项
2. 您可以使用Dask这样的工具将其扩展到单个机器的多个核心（也可以使用Numba）或跨集群
3. 您可以使用其中一个GPU阵列库，如Torch、TensorFlow、CuPy或Jax，在GPU上运行它

你也可以做任何上述的混合。你知道吗