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

Question

I am looping through a large number of H x W matrices. I cannot store them all in memory. I need to get N matrices. For example, the element of the 1st of N matrix in position (i, j) will be the largest among all elements in position (i, j) of all processed matrix matrices. For the second of the N matrix, the elements that are the second-largest will be taken, and so on.

Example.

Let N = 2. Then the 1st matrix will look like this.

And the second matrix is like this.

How to do such an operation inside a loop so as not to store all matrices in memory?

Answer 1

The comments suggested using the np.partition function. I replaced the use of numpy with cupy , which uses the GPU. And also added a buffer to sort less frequently.

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]

The solution does not work very quickly, but I have come to terms with this speed.