>我遇到了一个实际问题,其中更多的矢量化比更少的矢量化慢。将主要问题简化为以下玩具模型。 在下面的代码中,我使用两种不同的方法来实现相同的功能。f1
使用一个循环,f2
使用两个循环。天真地,我们会认为f1
比f2
快.但是f2
比f1
快。我不明白为什么会这样。
import numpy as np
def time_function(f, *args):
import time
tic = time.time()
f(*args)
toc = time.time()
return toc - tic
def f2(X, Y):
d = np.zeros((X.shape[0], Y.shape[0]))
for i in range(X.shape[0]):
for j in range(Y.shape[0]):
d[i,j] = np.sum(np.square(X[i] - Y[j]))
return d
def f1(X, Y):
d = np.zeros((X.shape[0], Y.shape[0]))
for i in range(X.shape[0]):
d[i] = np.sum(np.square(Y - X[i]), axis=1)
return d
X = np.ones((500, 3072))
Y = np.ones((5000, 3072))
time2 = time_function(f2,X,Y)
print('Two loop version took %f seconds' % time2)
time1 = time_function(f1,X,Y)
print('One loop version took %f seconds' % time1)
Two loop version took 24.691270 seconds
One loop version took 31.785896 seconds
我猜罪魁祸首隐藏在f1:
d[i] = np.sum(np.square(Y - X[i]), axis=1)
每次迭代都从整个 Y 数组中减去 X[i],这会导致密集广播,涉及 0 <= i <= 5000 范围内的迭代,而f2
减去两个不同的值,以0 <= i <= 500为界
在我的机器上,f1
稍微快一些。 但是完全矢量化的版本
D3=(np.square(Y[None,:,:]-X[:,None,:])).sum(axis=2)
给出内存错误,因为它必须创建 (500, 5000, 3072( 阵列 (57G(。
迭代和内存管理之间需要权衡。 通常,对相对复杂的任务进行几次迭代比需要分配更大矩阵的一个步骤更快。 在您的情况下,存在Y-X
差异,但是在用sum
减小尺寸之前,它还必须执行square
(相同的大尺寸(。
In [23]: timeit -n1 -r1 f2(X,Y)
1min 21s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)
In [24]: timeit -n1 -r1 f1(X,Y)
1min 13s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)
f1
的变体,在1min 25s
迭代 5000 行Y
(而不是 500 行X
(时间。 几次迭代通常比多次迭代好,前提是内存管理问题不会对您造成不利影响。
努巴
使用numba
的迭代但编译版本明显更好:
In [32]: import numba
...: @numba.jit(nopython=True)
...: def f3(X, Y):
...: d = np.zeros((X.shape[0], Y.shape[0]))
...: for i in range(X.shape[0]):
...: for j in range(Y.shape[0]):
...: d[i,j] = np.sum(np.square(X[i] - Y[j]))
...: return d
...:
In [33]: D3=f3(X,Y)
In [34]: np.allclose(D,D3)
Out[34]: True
In [35]: timeit -n1 -r1 f3(X,Y)
20 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)
并且使总和迭代显式可以节省更多时间:
In [36]: @numba.jit(nopython=True)
...: def f4(X, Y):
...: d = np.zeros((X.shape[0], Y.shape[0]))
...: for i in range(X.shape[0]):
...: for j in range(Y.shape[0]):
...: for k in range(X.shape[1]):
...: d[i,j] += (X[i,k] - Y[j,k])**2
...: return d
...:
...:
In [37]: D4 = f4(X,Y)
In [38]: np.allclose(D,D4)
Out[38]: True
In [39]: timeit -n1 -r1 f4(X,Y)
10.7 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)
我在两个定义中都看到了迭代。矢量化仅在大批量时获得收益!实际上,您一次只矢量化一行值的值。Python 循环能够处理这个问题,相比之下,进行和设置矢量化计算的开销需要很长时间。由于涉及迭代,f1
未完全矢量化。因此,当迭代就足够时,这种比较对矢量化是不公平的。
tl;DR 具有足够的值,矢量化"设置"被矢量化本身的增益所抵消。