利用 NumPy 广播机制提升机器学习运算性能

Python 能够力克群雄，成为科学计算、人工智能领域的最热语言，其数学工具包 NumPy 可谓居功至伟。但由于要兼顾建模能力和运算性能，NumPy 相当抽象，写出来的代码非常精简高效，令人拍案叫绝。我常常感觉能读懂 NumPy 代码就非常烧脑了，自己要写，只能是望洋兴叹吧。

 

计算相似度

 

这几天做一些数据试验，需要计算一批向量两两之间的相似度，例如下面这个矩阵（以下称其为 U），从第 0 行到第 5 行，每一行都可以看作一个 6 维线性空间的向量。

 

 

如果用欧几里得距离来度量向量之间的相似度（不要被复制粘贴过来的公式唬住，介只不过就是初中学的勾股定理）：

 

 

那么，这 6 个向量两两相似度的计算结果就是下面这样一个矩阵（以下称其为 D），例如其中第 1 行第 2 列的数值 2.40 就是第 1 个向量与第 2 个向量之间的欧几里得距离。当然，这个矩阵是对称的。

 

 

下面我们看看怎么实现这样的两两相似度计算。

 

用 Python 循环计算相似度

 

用最自然的姿势撸出来的代码是这样的：

import numpy as np

U = np.array(...这里是上文说的矩阵U的数据...)

r = range(6)

D = [

    [

        np.sqrt(((U[i,:]-U[j,:])**2).sum())

        for j in r

    ]

    for i in r

]

“[ ... for i in f ]”是 Python 的语法糖，就是精简化的 for 循环语句。我们使用 i、j 两重循环，依次从矩阵 U 中取出第 i、j 行向量（即 U[i,:] 和 U[j,:]），然后计算它们之间的欧几里得距离：np.sqrt(((U[i,:]-U[j,:])**2).sum())。计算结果 D 就是我们要的相似度矩阵。

 

在我的电脑上，用 iPython 的“%timeit” 语句将上述代码执行 1 万次，测出平均每次执行速度是 0.297 毫秒。

 

Jupyter Notebook 和 iPython 是 Python 的交互式开发环境，也是 Python 得以在科学家（包括数据科学家）圈子流行的主要功臣。

 

用 SciKit-Learn 计算相似度

 

计算相似度是比较常见的需求，肯定可以借助现成的工具来实现，比如 SciKit-Learn。SciKit-Learn 是非常流行的机器学习工具包，不过只是涵盖了传统的机器学习算法，也是基于 NumPy 的。

from sklearn.metrics.pairwise import euclidean_distances

D = euclidean_distances(U,U)

 

现成的工具用起来就是方便！由于是计算同一组向量相互间的距离，所以两个参数都是 U。

 

用“%timeit”测出这个语句的平均执行速度是 0.073 毫秒，比 Python 循环提升了整整 3 倍速度。

 

虽然没有去看 euclidean_distances 的源代码，我敢肯定不是用 Python 循环来实现的，一定是直接利用了 NumPy 的矩阵计算能力，否则不会这么快。

 

我仔细“端详”着 U 矩阵，想象着向量相互之间的交叉运算，感觉这就是传说中的 NumPy 广播（Broadcasting）的运算方式......

 

NumPy 广播

 

简单地理解，NumPy 的“广播（Broading）”机制就是：在作两个数组的对应元素间的计算（Element-wise Operation）时，如果两个数组在某个方向上的元素个数不相等，就把数据复制若干份，把元素个数凑成相等，就能执行对应元素的计算了。“广播”的机制非常抽象，讲再多也讲不清除，必须举几个具体的例子。

 

例一：单个数据在一个方向上“广播”：

 

 

例二：两个向量分别在两个不同的方向上“广播”： 

 

 

例三：两个矩阵分别在两个不同的方向（三维空间中的方向）上“广播”：

 

 

例三正是我们要利用的 NumPy 广播方式。“广播”过程是由 NumPy 自动完成的，不需要我们写什么代码。

 

用 NumPy 广播实现相似度计算

 

...... 我决定挑战一下，时间一分钟一分钟过去，终于拼出了如下代码：

 

D = np.sqrt(((U[:,np.newaxis,:] - U[np.newaxis,:,:])**2).sum(axis=2))

 

要利用广播机制，就得想办法把二维的矩阵 U 投射到高维空间，造成某个维度上元素个数的不相等。“np.newaxis”是 NumPy 的一个常量，用于创建一个新的维度。“U[:, np.newaxis, :]”和“U[np.newaxis, :, :]”分别将矩阵 U 这个二维平面投射到三维空间中两个不同的方向（见上面的“例三”）。这两个平面相减（减法是一种 Element-wise 计算）就会触发广播机制，从而得到一个三维的差值数组（即三维张量 Tensor）。“**2”是对这些差值求平方，“sum(axis=2)”是将平方值沿着第三维度加到一起，再经过“np.sqrt”开方就得到了正确的结果。

 

再用“%timeit”测一下，平均 0.013 毫秒，速度比 SciKit-Learn 提升了 5 倍，比 Python 循环提升了 20 倍以上。成功！

 

更快、更强

 

其实使用 NumPy 更大的好处在于：新一代的科学计算、人工智充工具包，例如 TensorFlow、PyTorch、Numba等，都继续了 NumPy 的接口，但可以通过 GPU 来完成具体的计算。如果有成千上万的数据点需要计算，GPU 的优势就能突显出来了，可以数十倍地提升运算速度。

 

如果数据量再大几个数量级，比如上亿用户、上百万产品项，那么单个 GPU 也搞不定了，需要分布式 GPU 集群才能完成数据模型的训练。现有的分布式集群例如 Tensorflow Serving、Caffe2，此外还有一大批基于“Spark + Kubernetes + GPU容器”的集群工具正在赶来......    欢迎来微信公众号《意庭斯言》找我！