numpy argsort 性能下降

Question

如果這個問題看起來很基本，我真誠地道歉。

鑒於：

>>> import numpy as np
>>> import time
>>> A = np.random.rand( int(1e5), int(5e4) ) # large numpy array

目標：

>>> bt=time.time(); B=np.argsort(A,axis=1);et=time.time();print(f"Took {(et-bt):.2f} s")

但是，計算索引數組需要很長時間：

# Took 316.90 s

題：

有沒有其他時間有效的方法來做到這一點？

干杯，

Answer 1

輸入數組A的形狀為(100_000, 50_000)並且默認包含np.float64值。 這意味着您需要8 * 100_000 * 50_000 / 1024**3 = 37.2 Gio的內存僅用於此數組。 您可能還需要為輸出矩陣B （它應該包含類型為np.int64項目）提供相同數量的空間。 這意味着您需要一台至少具有 74.4 Gio 的機器，更不用說操作系統 (OS) 和運行軟件所需的空間（因此可能至少為 80 Gio）。 如果您沒有這樣的內存空間，那么您的操作系統將使用您的存儲設備作為交換內存，這會慢得多。

假設你有這樣一個可用的內存空間，這樣的計算是非常昂貴的。 這主要是由於填充B數組時出現頁面錯誤，以及B非常大以及 Numpy 的默認實現按順序進行計算的事實。 您可以使用並行 Numba 代碼和較小的輸出來加快計算速度。 下面是一個例子：

import numba as nb

@nb.njit('int32[:,::1](float64[:,::1])', parallel=True)
def fastSort(a):
    b = np.empty(a.shape, dtype=np.int32)
    for i in nb.prange(a.shape[0]):
        b[i,:] = np.argsort(a[i,:])
    return b

請注意，Numba的實施argsort比與NumPy的之一，但其水貨版本低效率要快很多，如果目標機器有MNY核心和良好的內存帶寬。

以下是我的 6 核機器在大小矩陣(10_000, 50_000) （小 10 倍，因為我沒有 80 Gio 的 RAM）：

Original implementation: 28.14 s
Sequential Numba:        38.70 s
Parallel Numba:           6.79 s

因此，所得解決方案的速度提高了 4.1 倍。

請注意，在這種特定情況下，您甚至可以對B的項目使用類型uint16 ，因為每行的大小小於2**16 = 65536 。 這可能不會明顯更快，但它應該節省一些額外的內存。 由此產生的所需內存將為 46.5 Gio。 您可以使用np.float32類型進一步減少所需的內存量（通常以犧牲精度為代價）。

如果您想進一步縮短執行時間，那么您需要使用 C 或 C++ 等低級語言為您的特定需求實現更快的argsort實現。 但是請注意，如果您不是使用這種語言的有經驗的程序員或不熟悉低級優化，那么擊敗 Numpy 絕非易事。 如果您對這樣的解決方案感興趣，一個好的開始可能是實現基數排序。