如何有效地為多個矩陣中的每個元素分別找到 N 個最大值？

Question

我正在遍歷大量 H x W 矩陣。 我無法將它們全部存儲在 memory 中。 我需要得到 N 個矩陣。 例如，position (i, j) 中 N 矩陣的第一個元素將是所有已處理矩陣矩陣中 position (i, j) 中所有元素中最大的元素。 對於 N 矩陣中的第二個，將取第二大的元素，依此類推。

例子。

讓 N = 2。那么第一個矩陣將如下所示。

第二個矩陣是這樣的。

如何在循環內進行這樣的操作，以免將所有矩陣存儲在 memory 中？

Answer 1

評論建議使用np.partition function。 我用 cupy 替換了numpy的使用，它使用了 GPU。 並且還添加了一個緩沖區來減少排序頻率。

import cupy as np

buf = // # As much as fits into the GPU
largests = np.zeros((buf + N, h, w))
for i in range(num):
    val = //
    largests[i % buf] = val
    if i % buf == buf - 1:
        largests.partition(range(buf, buf + N), axis=0)
largests.partition(range(buf, buf + N), axis=0)  # Let's not forget the tail
res = largests[:-(N + 1):-1]

該解決方案的工作速度不是很快，但我已經接受了這種速度。