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

问题描述

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

示例。

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

第二个矩阵是这样的。

如何在循环内部进行这样的操作，以免将所有矩阵都存储在内存中？

解决方法

建议使用 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,axis=0)  # Let's not forget the tail
res = largests[:-(N + 1):-1]

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