如何在多个矩阵中有效地分别为每个元素找到 N 个最大值？

Question

我正在遍历大量 H x W 矩阵。我无法将它们全部存储在内存中。我需要得到 N 个矩阵。例如，1st of N matrix 在位置 (i, j) 的元素将是所有处理过的 matrix 矩阵中位置 (i, j) 的所有元素中最大的。对于N矩阵的第二个，取第二大的元素，以此类推。

示例。

令 N = 2。那么第一个矩阵将如下所示。

而第二个矩阵是这样的

如何在循环中做这样的操作，才不会把所有的矩阵都存储在内存中？

Answer 1

评论建议使用np.partition功能。我用使用 GPU 的 cupy 替换了 numpy 的使用。并且还添加了一个缓冲区来减少排序。

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]

该解决方案的工作速度不是很快，但我已经接受了这种速度。

如何在多个矩阵中有效地分别为每个元素找到 N 个最大值？

How to efficiently find separately for each element N maximum values among multiple matrices?

python

algorithm

numpy

matrix

memory-efficient