使用numpy进行for循环的多核处理

时间:2018-04-18 14:02:17

标签: python numpy multicore

我使用numpy计算了矢量。 如何使用多核和numpy计算向量?

import numpy as np

num_row, num_col = 6000, 13572

ss = np.ones((num_row, num_col), dtype=np.complex128)
ph = np.random.standard_normal(num_row)
fre = np.random.standard_normal(num_row)
tau = np.random.standard_normal(num_col)

for idx in range(num_row):
    ss[idx, :] *= np.exp(1j*(ph[idx] + fre[idx]*tau))

1 个答案:

答案 0 :(得分:4)

我们可以利用broadcasting来获得基于NumPy的解决方案 -

ss = np.exp(1j*(ph[:,None] + fre[:,None]*tau))

将此移植到numexpr以利用快速transcendental operations alongwith multi-core capability -

import numexpr as ne

def numexpr_soln(ph, fre):
    ph2D = ph[:,None]
    fre2D = fre[:,None]
    return ne.evaluate('exp(1j*(ph2D + fre2D*tau))')

计时 -

In [23]: num_row, num_col = 6000, 13572
    ...: ss = np.ones((num_row, num_col), dtype=np.complex128)
    ...: ph = np.random.standard_normal(num_row)
    ...: fre = np.random.standard_normal(num_row)
    ...: tau = np.random.standard_normal(num_col)

# Original soln
In [25]: %%timeit
    ...: for idx in range(num_row):
    ...:     ss[idx, :] *= np.exp(1j*(ph[idx] + fre[idx]*tau))
1 loop, best of 3: 4.46 s per loop

# Native NumPy broadcasting soln
In [26]: %timeit np.exp(1j*(ph[:,None] + fre[:,None]*tau))
1 loop, best of 3: 4.58 s per loop

对于具有不同内核/线程数的Numexpr解决方案 -

# Numexpr solution with # of threads = 2
In [51]: ne.set_num_threads(nthreads=2)
Out[51]: 2

In [52]: %timeit numexpr_soln(ph, fre)
1 loop, best of 3: 2.18 s per loop

# Numexpr solution with # of threads = 4
In [45]: ne.set_num_threads(nthreads=4)
Out[45]: 4

In [46]: %timeit numexpr_soln(ph, fre)
1 loop, best of 3: 1.62 s per loop

# Numexpr solution with # of threads = 8
In [48]: ne.set_num_threads(nthreads=8)
Out[48]: 8

In [49]: %timeit numexpr_soln(ph, fre)
1 loop, best of 3: 898 ms per loop