K-Nearest Neighbors (KNN) 是一个常用的分类算法,但是两两之间计算距离是比较耗费时间的,通过向量化(vectorize) 的方法,可以极大的加速计算的过程。
直接计算
最直接的算法是通过两个循环来计算每一组距离,代码如下:
import numpy as np
def brute_force(X_train, X):
num_test = X.shape[0]
num_train = X_train.shape[0]
dists = np.zeros((num_test, num_train))
for i in xrange(num_test):
for j in xrange(num_train):
dists[i, j] = np.linalg.norm(X[i] - X_train[j])
return dists
其中每一个数据是一个D
维的向量, X_train
是一个num_train * D
的矩阵,表示训练数据,X
是一个 num_test * D
的矩阵,表示测试数据,而要计算的距离为欧式距离。
去掉内层循环
通过 Numpy 的官方文档 可以发现,numpy.linalg.norm
支持 axis
参数,因此可以直接去掉内层的循环,将代码改写如下:
import numpy as np
def one_loop(X_train, X):
num_test = X.shape[0]
num_train = X_train.shape[0]
dists = np.zeros((num_test, num_train))
for i in xrange(num_test):
dists[i, :] = np.linalg.norm(X_train - X[i],axis=1)
return dists
去掉外循环:第一次尝试
去掉内循环的方法无法直接用于去掉外层的循环,这时候需要进一步开阔思路。既然我们要求的是每一个测试数据,每一个训练数据的组合在每一个维度上的差值,那么可以把数据映射到三维空间,diff [i, j, k]
表示第i
个测试数据与第j
个训练数据的组合在第k
维度的差别,然后再映射回两维空间,借助 Numpy 的 broadcast 机制,可以将代码改写为:
import numpy as np
def first_try(X_train, X):
num_test = X.shape[0]
num_train = X_train.shape[0]
diff = np.subtract(X[:,np.newaxis, :], X_train)
dists = np.sqrt(np.multiply(diff,diff).sum(2)).reshape([num_test,num_train])
return dists
然而运行起来会发现,当数据量比较大的时候产生了 Memoey Error
,因为这个方法开辟了过大的数组。
去掉外循环:第二次尝试
既然如此,能不能不开辟中间过程中的临时三维数组来计算呢?这就需要进一步观察一下要计算的公式。
\begin{equation}
\begin{split}
dists[i, j] &= \sqrt{\sum_k{(X[i, k] - X\_train[j, k] )^2} } \\
&= \sqrt{\sum_k{\left(X[i, k]^2 - 2 \cdot X[i, k] \cdot X\_train[j, k] +X\_train[j, k] ^ 2 \right) } } \\
&= \sqrt{\sum_k{X[i, k]^2} - 2\sum_k{ (X[i, k] \cdot X\_train[j, k])} + \sum_k{X\_train[j, k] ^ 2 } }
\end{split}
\end{equation}
可以发现,需要的结果可以分成三部分来计算,代码如下:
import numpy as np
def second_try(X_train, X):
num_test = X.shape[0]
num_train = X_train.shape[0]
A = np.square(X).sum(axis = 1)
B = np.dot(X, X_train.T)
C = np.square(X_train).sum(axis = 1)
dists = np.asarray(np.sqrt(np.matrix(A).T-2*B+C))
return dists
运行结果
向量化在性能改进上的作用是显著的。选用num_train=5000
, num_test=500
, D=3072
的一个数据,在型号为Intel i5-3210m
的 CPU 上进行测试,得到结果如下:
直接计算: 59.594000 秒 完全向量化计算: 0.808000 秒
减少了大约 98.7%
的计算时间。