在 Python 中計算兩個向量的所有點之間的距離：linalg、ase.geometry、並行化、GPU 等

Question

我有兩組 500,000 個 3D 坐標， pos1和pos2 。 我需要計算這兩組所有點之間的距離，然后給出一個 (500000,500000) 數組。 所以pos1[0]需要與pos2[0] , pos2[1] , ... pos2[499999]進行比較，同樣適用於pos1[1] , ...

這是一個非常緩慢的過程，並且希望在具有 8 個節點的集群上少於一個小時。

我一直在嘗試以下 15,000 個職位，以了解如何進行：

• numpy.linalg.norm (43s): dist = np.linalg.norm(pos1[:, np.newaxis, :] - pos2, axis=2)
• numpy.einsum (12s):

代碼如下：

diff = pos1[:,None,:] - pos2
dist = np.sqrt(np.einsum('ijk,ijk->ij', diff, diff))

• ase.geometry.find_mic (3min36s):

代碼如下所示：

diff = pos2[:, np.newaxis, :]-list1
dist = np.empty((pos1.shape[0], pos2.shape[0]))
for i, d in enumerate(diff):
    print(i)
    dist[i, i:] = find_mic(d[i:], cell=[[50., 0.0, 0.0], [25., 45., 0.0], [0.0, 0.0, 100.]])[1]
    if i+1 < len(diff):
        dist[i+1:, i] = dist[i, i+1:].T

所以我的問題是：如何讓它工作得更快？ 由於 einsum 像 N^2 一樣縮放，因此需要 12*(500000/15000)^2/3600 = 3h42min 才能找到解決方案。

是否有其他（更快，或擴展性更好）的方法？ 我可以利用可以在 8 個節點上使用 20 個線程這一事實嗎？ 點 x 和點 y 之間的距離是否等於 y 到 x 的距離？ 我還可以訪問 GPU 節點，如果這樣可以更好的話。

Answer 1

這是一個立方體的例子。

您可以將立方體分成 32³ 更小的立方體，並計算每個立方體中心之間的距離 d（和 1/d）。 它節省了大量的計算，因為 32² 立方體共享相同的 x，32² 共享相同的 y，並且 32² 共享相同的 z，並且由於對稱性，您不需要計算所有立方體。 如果你計算 1/8 個立方體的距離，你就知道所有的距離。

假設你有一百萬分。

現在，每個立方體平均有 1,000,000/32³=30 個點

你需要計算這 30 個點到他的立方體中心的距離：(dx,dy,dz)。

點之間的距離就是立方體之間的距離d，更多的是修正dd=(x.dx+y.dy+z.dz)/d

[您已經預先計算了 (x,y,z)/d，並且對於立方體中的所有點都是相同的]

（我沒有做精確的計算。我只是給出大致的想法）

您仍然需要將每個點與每個點進行比較，即 O(n²)，但是您保存了每對點的根計算。

您可以將結果與精確計算進行比較，如果差異不夠小，您可以增加立方體的數量。

遠方立方體的精度會更高，因此您可能需要調整什么距離才能更好地計算所有內容，以及什么距離足以使用此技巧。

Answer 2

請重新考慮您的需求。 一個(500000,500000)的 64 位值數組大約需要 1862 GiB 的 RAM，這是非常巨大的，大多數計算機都沒有。 如果您最終將其存儲在 RAM 中，那么在 GPU 上計算它是沒有用的，因為計算是內存堆積並且CPU-GPU 傳輸比主 RAM 慢得多。 AFAIK，世界上還沒有 GPU 擁有約 2 TiB 的嵌入式 RAM。 將這個巨大的數組存儲在 RAM 中實際上是主要問題。 與緩存相比，主 RAM 很慢，CPU 已經可以非常快速地並行計算。 同樣的事情也適用於 GPU：VRAM 通常明顯更快，GPU 通常也更快（盡管 VRAM 和緩存更小）。 這意味着您應該即時計算數據，以便有效地使用現代計算機。 因此，如果您打算使用 GPU，那么使用矩陣的計算也應該在 GPU 上實現。 老實說，並行計算應該已經非常適合您的任務，我建議您僅在最快的 CPU 實現未能達到您的時序預算時才嘗試使用 GPU（因為有效地使用 GPU 比 CPU 困難得多）。 為了即時計算操作，@Colim 答案提供了一些實現該操作的提示（在 CPU 和 GPU 上）。

這個問題是一個XY 問題：實際問題不是計算距離矩陣，而是你想用它做什么以及如何有效地完成它（計算整個矩陣實際上是一個低效的解決方案）。 該方法的復雜度為O(n^2)並且無法改進，因為輸出的大小為O(n^2) 。 生成距離矩陣很少有用，並且已知不能在大型數據集上進行縮放。 有相對快速的算法可以對大型數據集進行分類、計算最近鄰等，而無需計算整個矩陣（通常在O(n log n)時間內）。