矩陣之間的行向量兩兩距離
給定數據矩陣\(A\in R^{n\times d}\)和矩陣\(B\in R^{m\times d}\) ,A和B中每一行都是一個數據點,現在要去求A中所有元素和B中所有元素之間的歐氏距離。即計算矩陣\(D =(d_{ij} = ||a_{i,:} - b_{j, :}||) \in R^{n\times m}\)。
直接去做\(n\times m\)的兩層循環當然可以得到結果,但是既然數據都已經放在矩陣里了,它自然可以“並行”地去計算,好利用一些現有的矩陣計算庫(比如numpy
、torch
)等的優化。
具體來講,把歐氏距離的運算展開
觀察可以知道,式子中的三項分別等於A中第i行的平方和、B中第j行的平方和、A中第i行和B中第j行的內積,這樣的話,可以分別對A和B做逐元素平方求和,而向量兩兩做內積,得到的結果其實就是A與B轉置相乘。矩陣相乘作為一個經典操作那是被優化得好好的,比起自己去做兩層循環方便得多。
整個過程就成了\(RowSum(A .* A)^T + RowSum(B .* B) - 2A\cdot B^T\) 。
這里有個問題,A平方和之后轉置得到的是\(n\times 1\) 的向量,B平方和得到的是\(1\times m\)的向量,要把它們都擴展成\(n\times m\)的向量才好相加。幸運的是,現在常見的矩陣運算庫基本都是有“廣播”功能的。當向量\(v \in R^{n\times 1}\)和向量\(u\in R^{1\times m}\) 相加時,它們分別會被自動擴展成為矩陣\(V\in R^{n\times m}\)和\(U\in R^{n\times m}\) ,顯然可以想到V就是把v復制了m列,U是把u復制了n行。通過自動廣播,維度的處理就完全不是問題了。
具體代碼如下
# numpy version
A2 = (A ** 2).sum(axis=1).reshape((-1, 1)) # 轉為n*1
B2 = (B ** 2).sum(axis=1) # 這里不用reshape了,(n,)和(1,n)的向量在計算過程中效果是一樣的
D = A2 + B2 - 2 * A.dot(B.transpose())
# D = np.clamp(D, a_min=0) # 有時候會出現精度問題導致結果小於0 ,我習慣於在這里進行一次校正
D = D.sqrt() # 開方得到歐氏距離
# ...
# torch version
A2 = (A ** 2).sum(dim=1).reshape((-1, 1))
B2 = (B ** 2).sum(dim=1)
D = A2 + B2 - 2 * A.mm(B.t())
# ...
實際場景經常會遇到一個點矩陣\(A\in R^{a\times d}\) 需要求所有數據點之間的兩兩距離,此時就可以用上面的方式去計算A和A自身的兩兩距離。
批量計算點集與點集之間的兩兩距離
今天在做實驗時遇到這樣一個場景,\(X\in R^{B\times n\times d}\) 和\(Y\in R^{B\times m\times d}\),其實這也是玩深度學習時常會出現到的情況,就是數據集被划分為一個一個的Batch,這里我想對每一個batch b分別計算\(X[b,:, :]\) 和\(Y[b,:, :]\) 這兩個矩陣中的兩兩距離,放在常用的計算庫(如torch
、paddle
等)中,一樣是可以轉換為矩陣運算的,只是需要做一點點的reshape而已。
# paddle version
X2 = (X ** 2).sum(axis=2).reshape((batch_size, 1, -1))
Y2 = (Y ** 2).sum(axis=2).reshape((batch_size, -1, 1)) # 由於多了第一維,這里的reshape不能偷懶不寫了
D = X2 + Y2 - 2 * X.bmm(Y.transpose((0, 2, 1)))
之所以能這么做,得益於很多庫實現了bmm
運算,即對每個batch分別進行矩陣乘法(這個函數應用於張量的時代,在numpy里還沒有)。
這一需求其實在pytorch中有一個特定的函數torch.cdist
完成了。然而很不幸的百度paddle框架中並沒有類似函數,所以只能自己實現了一版。當然torch內部似乎是用了contiguous函數去做了些矩陣結構的調整(從GitHub issue中看到),具體實現細節倒沒有去看。