这是上一个问题的延续 数据是
floatMatrix
,其维度为 (750000, 4)。函数 testFunction[x,y,w,z]
是 4 个变量的多项式函数,返回 4D 向量,适用于所有 750000 个向量。
对于 Mathematica,我使用可列出的
Compile
和 Parallelization->False
。
Compile[{{f, _Real, 1}},
{
{0.011904761904761973` f[[2]]f[[1]]^3 +
0.002976190476190474` f[[1]]f[[2]]^3 - 0.020833333333333325` f[[3]] +
0.002976190476190474` f[[3]]^3 +
f[[2]]^2 (0.0029761904761904778` f[[3]] +...
{0.002976190476190483` f[[1]]^3 + 0.011904761904761906` f[[2]]^3 -
0.0875` f[[3]] + 0.0029761904761904765` f[[3]]^3 +
f[[1]]^2 (0.005952380952380952` f[[2]] +...
},CompilationTarget -> "C", RuntimeAttributes -> {Listable},
Parallelization -> True];
time = RepeatedTiming[testFunction[floatMatrix]];
Print["In Mathematica-C it takes an average of ", time[[1]], " secs."]
对于 Python,我使用 NumPy。我有两个实现:
实现1(内部转置):
def testFunction(data):
f1, f2, f3, f4 = data.T
results = np.zeros((data.shape[0], 4)) # Initialize a results array
results[:, 0] = (0.011904761904761973*f2*f1**3 + 0.002976190476190474*f1*f2**3 -
0.020833333333333325*f3 + 0.002976190476190474*f3**3 + f2**2*
(0.0029761904761904778*f3 +...
results[:, 1] = (0.002976190476190483*f1**3 + 0.011904761904761906*f2**3 - 0.0875*f3
+ 0.0029761904761904765*f3**3 + f1**2*(0.005952380952380952*f2 +
0.002976190476190469*f3 + 0.0029761904761904726*f4) +...
return results
duration=0
for i in range(10):
start_time = time.time()
testFunction(floatMatrix)
end_time = time.time()
duration = duration + end_time - start_time
duration=duration*0.1
print(f"With numpy it takes an average of': {duration} seconds")
实现2(外部转置):
def testFunction(f1, f2, f3, f4):
results = np.zeros((f1.shape[0], 4)) # Initialize a results array
results[:, 0] = ...
results[:, 1] = ...
results[:, 2] = ...
results[:, 3] = ...
return results
# Transpose the data outside the function
f1, f2, f3, f4 = floatMatrix.T
duration=0
for i in range(10):
start_time = time.time()
testFunction(f1, f2, f3, f4)
end_time = time.time()
duration = duration + end_time - start_time
duration=duration*0.1
print(f"With numpy vectorized operations it takes an average of': {duration} seconds")
这些是时间:
我原以为 NumPy 会快得多。我可以对 Python 中的代码进行任何更改,使其比 Mathematica 编译的函数更快吗?
有多种方法可以使 NumPy 代码更快。一种方法是打开 Numba,这是一种优化器,它使用编译器来加速 NumPy 繁重代码。
import numba as nb
@nb.njit()
def test_function():
...
这样做,我获得了 6 倍的加速。
接下来,我尝试打开fastmath。这种模式允许 Numba 改变某些浮点运算以提高速度。 (例如,通常不允许将 (a+b)+c 重写为 a+(b+c),但在 fastmath 下这是允许的转换。)
import numba as nb
@nb.njit(fastmath=True):
def test_function()
...
这可提高 3% 的速度。
接下来,我注意到,相对于 Mathematica 执行计算的顺序,执行此计算的顺序具有较差的内存局部性。 Mathematica 一次计算一个向量,但 NumPy 一次计算表达式的一个组成部分,这需要它多次将
f1
列加载到内存中。
为了解决这个问题,我引入了一个显式循环。 (有关详细信息,请参阅代码部分。)这使我的速度又提高了 44%。
结合这些优化,与原始 Python 实现相比,我获得了 9.5 倍的加速。 (我没有对 Mathematica 代码进行基准测试,因为我没有副本。)
原始(加上修复语法错误的修复)
floatMatrix = np.random.rand(1200000, 4)
def test_function_transpose_inside(data):
f1, f2, f3, f4 = data.T
results = np.zeros((data.shape[0], 4)) # Initialize a results array
results[:, 0] = (0.011904761904761973*f2*f1**3 + 0.002976190476190474*f1*f2**3 -
0.020833333333333325*f3 + 0.002976190476190474*f3**3 + f2**2*
(0.0029761904761904778*f3))
results[:, 1] = (0.002976190476190483*f1**3 + 0.011904761904761906*f2**3 - 0.0875*f3
+ 0.0029761904761904765*f3**3 + f1**2*(0.005952380952380952*f2 +
0.002976190476190469*f3 + 0.0029761904761904726*f4))
return results
与 Numba
import numba as nb
@nb.njit()
def test_function_numba(data):
f1, f2, f3, f4 = data.T
results = np.zeros((data.shape[0], 4)) # Initialize a results array
results[:, 0] = (0.011904761904761973*f2*f1**3 + 0.002976190476190474*f1*f2**3 -
0.020833333333333325*f3 + 0.002976190476190474*f3**3 + f2**2*
(0.0029761904761904778*f3))
results[:, 1] = (0.002976190476190483*f1**3 + 0.011904761904761906*f2**3 - 0.0875*f3
+ 0.0029761904761904765*f3**3 + f1**2*(0.005952380952380952*f2 +
0.002976190476190469*f3 + 0.0029761904761904726*f4))
return results
Numba 和快速数学
import numba as nb
@nb.njit(fastmath=True)
def test_function_numba_fastmath(data):
f1, f2, f3, f4 = data.T
results = np.zeros((data.shape[0], 4)) # Initialize a results array
results[:, 0] = (0.011904761904761973*f2*f1**3 + 0.002976190476190474*f1*f2**3 -
0.020833333333333325*f3 + 0.002976190476190474*f3**3 + f2**2*
(0.0029761904761904778*f3))
results[:, 1] = (0.002976190476190483*f1**3 + 0.011904761904761906*f2**3 - 0.0875*f3
+ 0.0029761904761904765*f3**3 + f1**2*(0.005952380952380952*f2 +
0.002976190476190469*f3 + 0.0029761904761904726*f4))
return results
Numba、快速数学和循环
import numba as nb
@nb.njit(fastmath=True)
def test_function_numba_loop(data):
results = np.zeros((data.shape[0], 4)) # Initialize a results array
for i in range(data.shape[0]):
f1, f2, f3, f4 = data[i]
results[i, 0] = (0.011904761904761973*f2*f1**3 + 0.002976190476190474*f1*f2**3 -
0.020833333333333325*f3 + 0.002976190476190474*f3**3 + f2**2*
(0.0029761904761904778*f3))
results[i, 1] = (0.002976190476190483*f1**3 + 0.011904761904761906*f2**3 - 0.0875*f3
+ 0.0029761904761904765*f3**3 + f1**2*(0.005952380952380952*f2 +
0.002976190476190469*f3 + 0.0029761904761904726*f4))
return results
%timeit test_function_transpose_inside(floatMatrix)
215 ms ± 14.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit test_function_numba(floatMatrix)
33.6 ms ± 344 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit test_function_numba_fastmath(floatMatrix)
32.6 ms ± 452 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit test_function_numba_loop(floatMatrix)
22.5 ms ± 68.6 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)