问题描述
设 ${X_{m \times k}} = \left[ {\vec x_1^T;\vec x_2^T; \cdots ;\vec x_m^T} \right]$ (; 表示纵向连接) 和 ${Y_{n \times k}} = \left[ {\vec y_1^T;\vec y_2^T; \cdots ;\vec y_n^T} \right]$ , 计算矩阵 ${X_{m \times k}}$ 中每一个行向量和矩阵 ${Y_{n \times k}}$ 中每一个行向量的平方欧氏距离 (pairwise squared Euclidean distance), 即计算:
${\begin{bmatrix} {\left\| {{{\vec x}_1} - {{\vec y}_1}} \right\|_2^2}&{\left\| {{{\vec x}_1} - {{\vec y}_2}} \right\|_2^2}& \cdots &{\left\| {{{\vec x}_1} - {{\vec y}_n}} \right\|_2^2} \\ {\left\| {{{\vec x}_2} - {{\vec y}_1}} \right\|_2^2}&{\left\| {{{\vec x}_2} - {{\vec y}_2}} \right\|_2^2}& \cdots &{\left\| {{{\vec x}_2} - {{\vec y}_n}} \right\|_2^2} \\ \vdots & \vdots & \ddots & \vdots \\ {\left\| {{{\vec x}_m} - {{\vec y}_1}} \right\|_2^2}&{\left\| {{{\vec x}_m} - {{\vec y}_2}} \right\|_2^2}& \cdots &{\left\| {{{\vec x}_m} - {{\vec y}_n}} \right\|_2^2} \end{bmatrix}}$ (这是一个 $m \times n$ 矩阵).
这个计算在度量学习, 图像检索, 行人重识别等算法的性能评估中有着广泛的应用.
公式转化
在 NumPy 中直接利用上述原式来计算两个矩阵的成对平方欧氏距离, 要显式地使用二重循环, 而在 Python 中循环的效率是相当低下的. 如果想提高计算效率, 最好是利用 NumPy 的特性将原式转化为数组/矩阵运算. 下面就尝试进行这种转化.
先将原式展开为:
$$ \begin{bmatrix} {\left\| {{{\vec x}_1}} \right\|_2^2}&{\left\| {{{\vec x}_1}} \right\|_2^2}& \cdots &{\left\| {{{\vec x}_1}} \right\|_2^2} \\ {\left\| {{{\vec x}_2}} \right\|_2^2}&{\left\| {{{\vec x}_2}} \right\|_2^2}& \cdots &{\left\| {{{\vec x}_2}} \right\|_2^2} \\ \vdots & \vdots & \ddots & \vdots \\ {\left\| {{{\vec x}_m}} \right\|_2^2}&{\left\| {{{\vec x}_m}} \right\|_2^2}& \cdots &{\left\| {{{\vec x}_m}} \right\|_2^2} \end{bmatrix} + {\begin{bmatrix} {\left\| {{{\vec y}_1}} \right\|_2^2}&{\left\| {{{\vec y}_2}} \right\|_2^2}& \cdots &{\left\| {{{\vec y}_n}} \right\|_2^2} \\ {\left\| {{{\vec y}_1}} \right\|_2^2}&{\left\| {{{\vec y}_2}} \right\|_2^2}& \cdots &{\left\| {{{\vec y}_n}} \right\|_2^2} \\ \vdots & \vdots & \ddots & \vdots \\ {\left\| {{{\vec y}_1}} \right\|_2^2}&{\left\| {{{\vec y}_2}} \right\|_2^2}& \cdots &{\left\| {{{\vec y}_n}} \right\|_2^2} \end{bmatrix}} - 2 {\begin{bmatrix} {\left\langle {{{\vec x}_1},{{\vec y}_1}} \right\rangle }&{\left\langle {{{\vec x}_1},{{\vec y}_2}} \right\rangle }& \cdots &{\left\langle {{{\vec x}_1},{{\vec y}_n}} \right\rangle } \\ {\left\langle {{{\vec x}_2},{{\vec y}_1}} \right\rangle }&{\left\langle {{{\vec x}_2},{{\vec y}_2}} \right\rangle }& \cdots &{\left\langle {{{\vec x}_2},{{\vec y}_n}} \right\rangle } \\ \vdots & \vdots & \ddots & \vdots \\ {\left\langle {{{\vec x}_m},{{\vec y}_1}} \right\rangle }&{\left\langle {{{\vec x}_m},{{\vec y}_2}} \right\rangle }& \cdots &{\left\langle {{{\vec x}_m},{{\vec y}_n}} \right\rangle } \end{bmatrix}} $$
下面逐项地化简或转化为数组/矩阵运算的形式:
$$\left[ {\begin{matrix} {\left\| {{{\vec x}_1}} \right\|_2^2}&{\left\| {{{\vec x}_1}} \right\|_2^2}& \cdots &{\left\| {{{\vec x}_1}} \right\|_2^2} \\ {\left\| {{{\vec x}_2}} \right\|_2^2}&{\left\| {{{\vec x}_2}} \right\|_2^2}& \cdots &{\left\| {{{\vec x}_2}} \right\|_2^2} \\ \vdots & \vdots & \ddots & \vdots \\ {\left\| {{{\vec x}_m}} \right\|_2^2}&{\left\| {{{\vec x}_m}} \right\|_2^2}& \cdots &{\left\| {{{\vec x}_m}} \right\|_2^2} \end{matrix}} \right] = \left[ {\begin{matrix} {\left\| {{{\vec x}_1}} \right\|_2^2} \\ {\left\| {{{\vec x}_2}} \right\|_2^2} \\ \vdots \\ {\left\| {{{\vec x}_m}} \right\|_2^2} \end{matrix}} \right]\vec 1_n^T = \left( {\left( {X \circ X} \right){{\vec 1}_k}} \right)\vec 1_n^T = \left( {X \circ X} \right){\vec 1_k}\vec 1_n^T$$
式中, $\circ$ 表示按元素积 (element-wise product), 又称为 Hadamard 积; ${\vec 1_k}$ 表示维的全1向量 (all-ones vector), 余者类推. 上式中 ${\vec 1_k}$ 的作用是计算 $X \circ X$ 每行元素的和, 返回一个列向量; $\vec 1_n^T$ 的作用类似于 NumPy 中的广播机制, 在这里是将一个列向量扩展为一个矩阵, 矩阵的每一列都是相同的.
$$\left[ {\begin{matrix} {\left\| {{{\vec y}_1}} \right\|_2^2}&{\left\| {{{\vec y}_2}} \right\|_2^2}& \cdots &{\left\| {{{\vec y}_n}} \right\|_2^2} \\ {\left\| {{{\vec y}_1}} \right\|_2^2}&{\left\| {{{\vec y}_2}} \right\|_2^2}& \cdots &{\left\| {{{\vec y}_n}} \right\|_2^2} \\ \vdots & \vdots & \ddots & \vdots \\ {\left\| {{{\vec y}_1}} \right\|_2^2}&{\left\| {{{\vec y}_2}} \right\|_2^2}& \cdots &{\left\| {{{\vec y}_n}} \right\|_2^2} \end{matrix}} \right] = {\vec 1_m}{\left[ {\begin{matrix} {\left\| {{{\vec y}_1}} \right\|_2^2} \\ {\left\| {{{\vec y}_2}} \right\|_2^2} \\ \vdots \\ {\left\| {{{\vec y}_n}} \right\|_2^2} \end{matrix}} \right]^T} = {\vec 1_m}{\left( {\left( {Y \circ Y} \right){{\vec 1}_k}} \right)^T} = {\vec 1_m}\vec 1_k^T{\left( {Y \circ Y} \right)^T}$$
$$\left[ {\begin{matrix} {\left\langle {{{\vec x}_1},{{\vec y}_1}} \right\rangle }&{\left\langle {{{\vec x}_1},{{\vec y}_2}} \right\rangle }& \cdots &{\left\langle {{{\vec x}_1},{{\vec y}_n}} \right\rangle } \\ {\left\langle {{{\vec x}_2},{{\vec y}_1}} \right\rangle }&{\left\langle {{{\vec x}_2},{{\vec y}_2}} \right\rangle }& \cdots &{\left\langle {{{\vec x}_2},{{\vec y}_n}} \right\rangle } \\ \vdots & \vdots & \ddots & \vdots \\ {\left\langle {{{\vec x}_m},{{\vec y}_1}} \right\rangle }&{\left\langle {{{\vec x}_m},{{\vec y}_2}} \right\rangle }& \cdots &{\left\langle {{{\vec x}_m},{{\vec y}_n}} \right\rangle } \end{matrix}} \right] = \left[ {\begin{matrix} {\vec x_1^T} \\ {\vec x_2^T} \\ \vdots \\ {\vec x_m^T} \end{matrix}} \right]\left[ {\begin{matrix} {{{\vec y}_1}}&{{{\vec y}_2}}& \cdots &{{{\vec y}_n}} \end{matrix}} \right] = X{Y^T}$$
所以:
$$\left[ {\begin{matrix} {\left\| {{{\vec x}_1} - {{\vec y}_1}} \right\|_2^2}&{\left\| {{{\vec x}_1} - {{\vec y}_2}} \right\|_2^2}& \cdots &{\left\| {{{\vec x}_1} - {{\vec y}_n}} \right\|_2^2} \\ {\left\| {{{\vec x}_2} - {{\vec y}_1}} \right\|_2^2}&{\left\| {{{\vec x}_2} - {{\vec y}_2}} \right\|_2^2}& \cdots &{\left\| {{{\vec x}_2} - {{\vec y}_n}} \right\|_2^2} \\ \vdots & \vdots & \ddots & \vdots \\ {\left\| {{{\vec x}_m} - {{\vec y}_1}} \right\|_2^2}&{\left\| {{{\vec x}_m} - {{\vec y}_2}} \right\|_2^2}& \cdots &{\left\| {{{\vec x}_m} - {{\vec y}_n}} \right\|_2^2} \end{matrix}} \right] = \left( {X \circ X} \right){\vec 1_k}\vec 1_n^T + {\vec 1_m}\vec 1_k^T{\left( {Y \circ Y} \right)^T} - 2X{Y^T}$$
上述转化式中出现了 $X{Y^T}$ (矩阵乘). 矩阵乘在 NumPy 等很多库中都有高效的实现, 出现矩阵乘对代码的优化是有好处的.
特别地, 当 $X=Y$ 时, 原式等于 $\left( {X \circ X} \right){\vec 1_k}\vec 1_m^T + {\vec 1_m}\vec 1_k^T{\left( {X \circ X} \right)^T} - 2X{X^T}$ , 注意到第一项和第二项互为转置. 当 $\left( {X \circ X} \right){\vec 1_k} =\vec 1_m$ 且 $\left( {Y \circ Y} \right){\vec 1_k} =\vec 1_n$ (即 $X$ 和 $Y$ 的每一个行向量的范数均为1时), 原式等于 $2\vec 1_m \vec 1_n^T - 2X{Y^T} = 2\left(\vec 1_m \vec 1_n^T -X{Y^T}\right)$, $\vec 1_m \vec 1_n^T$ 是 $m \times n$ 全1矩阵.
代码实现
sklearn 中已经包含了用 NumPy 实现的计算 "两个矩阵的成对平方欧氏距离" 的函数 (sklearn.metrics.euclidean_distances [1] ), 它利用的就是上面的转化公式. 这里, 我们利用上面的转化公式并借鉴 sklearn, 用 NumPy 重新实现一个轻量级且易于理解的版本:
import numpy as np
def euclidean_distances(x, y, squared=True):
"""Compute pairwise (squared) Euclidean distances.
"""
assert isinstance(x, np.ndarray) and x.ndim == 2
assert isinstance(y, np.ndarray) and y.ndim == 2
assert x.shape[1] == y.shape[1]
x_square = np.sum(x*x, axis=1, keepdims=True)
if x is y:
y_square = x_square.T
else:
y_square = np.sum(y*y, axis=1, keepdims=True).T
distances = np.dot(x, y.T)
distances *= -2
distances += x_square
distances += y_square
np.maximum(distances, 0, distances)
if x is y:
distances.flat[::distances.shape[0] + 1] = 0.0
if not squared:
np.sqrt(distances, distances)
return distances
如果想进一步加速, 可以将
x_square = np.sum(x*x, axis=1, keepdims=True)
替换为
x_square = np.expand_dims(np.einsum('ij,ij->i', x, x), axis=1)
以及将
y_square = np.sum(y*y, axis=1, keepdims=True).T
替换为
y_square = np.expand_dims(np.einsum('ij,ij->i', y, y), axis=0)
使用 np.einsum 的好处是不会产生一个和 x 或 y 同样形状的临时数组 ( x*x 或 y*y 会产生一个和 x 或 y 同样形状的临时数组).
PyTorch 中也包含了计算 "两个矩阵的成对平方欧氏距离" 的函数 [2] , 不过它利用了如下的转化公式, 感兴趣的朋友可以自己用 NumPy 实现一下.
$$\begin{aligned} \left( {X \circ X} \right){{\vec 1}_k}\vec 1_n^T + {{\vec 1}_m}\vec 1_k^T{\left( {Y \circ Y} \right)^T} - 2X{Y^T} &= \left[ {\begin{matrix} { - 2X}&{\left( {X \circ X} \right){{\vec 1}_k}}&{{{\vec 1}_m}} \end{matrix}} \right]\left[ {\begin{matrix} {{Y^T}} \\ {\vec 1_n^T} \\ {\vec 1_k^T{{\left( {Y \circ Y} \right)}^T}} \end{matrix}} \right] \\ &= \left[ {\begin{matrix} { - 2X}&{\left( {X \circ X} \right){{\vec 1}_k}}&{{{\vec 1}_m}} \end{matrix}} \right]{\left[ {\begin{matrix} Y&{{{\vec 1}_n}}&{\left( {Y \circ Y} \right){{\vec 1}_k}} \end{matrix}} \right]^T} \\ \end{aligned}$$
另外上述的转化公式也可以用在其他 Python 框架 (如 TensorFlow) 或其他语言中, 这里就不展开叙述了.
参考
更新记录
- 20200605, 发布