NumPy之計算兩個矩陣的成對平方歐氏距離


問題描述

\({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), 即計算:

\(\left[ {\begin{array}{*{20}{c}} {\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{array}} \right]\) (這是一個 \(m \times n\) 矩陣).

這個計算在度量學習, 圖像檢索, 行人重識別等算法的性能評估中有着廣泛的應用.

公式轉化

在 NumPy 中直接利用上述原式來計算兩個矩陣的成對平方歐氏距離, 要顯式地使用二重循環, 而在 Python 中循環的效率是相當低下的. 如果想提高計算效率, 最好是利用 NumPy 的特性將原式轉化為數組/矩陣運算. 下面就嘗試進行這種轉化.

先將原式展開為:
\(\left[ {\begin{array}{*{20}{c}} {\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{array}} \right] + \left[ {\begin{array}{*{20}{c}} {\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{array}} \right] - 2\left[ {\begin{array}{*{20}{c}} {\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{array}} \right]\)

下面逐項地化簡或轉化為數組/矩陣運算的形式:
\(\left[ {\begin{array}{*{20}{c}} {\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{array}} \right] = \left[ {\begin{array}{*{20}{c}} {\left\| {{{\vec x}_1}} \right\|_2^2} \\ {\left\| {{{\vec x}_2}} \right\|_2^2} \\ \vdots \\ {\left\| {{{\vec x}_m}} \right\|_2^2} \end{array}} \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{array}{*{20}{c}} {\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{array}} \right] = {\vec 1_m}{\left[ {\begin{array}{*{20}{c}} {\left\| {{{\vec y}_1}} \right\|_2^2} \\ {\left\| {{{\vec y}_2}} \right\|_2^2} \\ \vdots \\ {\left\| {{{\vec y}_n}} \right\|_2^2} \end{array}} \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{array}{*{20}{c}} {\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{array}} \right] = \left[ {\begin{array}{*{20}{c}} {\vec x_1^T} \\ {\vec x_2^T} \\ \vdots \\ {\vec x_m^T} \end{array}} \right]\left[ {\begin{array}{*{20}{c}} {{{\vec y}_1}}&{{{\vec y}_2}}& \cdots &{{{\vec y}_n}} \end{array}} \right] = X{Y^T}\)

所以:

\(\left[ {\begin{array}{*{20}{c}} {\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{array}} \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), 它利用的就是上面的轉化公式. 這里, 我們利用上面的轉化公式並借鑒 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)
    # use inplace operation to accelerate
    distances *= -2
    distances += x_square
    distances += y_square
    # result maybe less than 0 due to floating point rounding errors.
    np.maximum(distances, 0, distances)
    if x is y:
        # Ensure that distances between vectors and themselves are set to 0.0.
        # This may not be the case due to floating point rounding errors.
        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*xy*y 會產生一個和 x 或 y 同樣形狀的臨時數組).

PyTorch 中也包含了計算 "兩個矩陣的成對平方歐氏距離" 的函數, 不過它利用了如下的轉化公式, 感興趣的朋友可以自己用 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{array}{*{20}{c}} { - 2X}&{\left( {X \circ X} \right){{\vec 1}_k}}&{{{\vec 1}_m}} \end{array}} \right]\left[ {\begin{array}{*{20}{c}} {{Y^T}} \\ {\vec 1_n^T} \\ {\vec 1_k^T{{\left( {Y \circ Y} \right)}^T}} \end{array}} \right] \\ &= \left[ {\begin{array}{*{20}{c}} { - 2X}&{\left( {X \circ X} \right){{\vec 1}_k}}&{{{\vec 1}_m}} \end{array}} \right]{\left[ {\begin{array}{*{20}{c}} Y&{{{\vec 1}_n}}&{\left( {Y \circ Y} \right){{\vec 1}_k}} \end{array}} \right]^T} \\ \end{aligned}\)

另外上述的轉化公式也可以用在其他 Python 框架 (如 TensorFlow) 或其他語言中, 這里就不展開敘述了.

參考

版權聲明

版權聲明:自由分享,保持署名-非商業用途-非衍生,知識共享3.0協議。
如果你對本文有疑問或建議,歡迎留言!轉載請保留版權聲明!


免責聲明!

本站轉載的文章為個人學習借鑒使用,本站對版權不負任何法律責任。如果侵犯了您的隱私權益,請聯系本站郵箱yoyou2525@163.com刪除。



 
粵ICP備18138465號   © 2018-2025 CODEPRJ.COM