1

How can I implement the tensorflow function tf.nn.top_k with Numpy? Suppose the input is ndarray in format heigh x width x channel?

neouyghur
  • 1,577
  • 16
  • 31

2 Answers2

2

You can use the answer here with Numpy 1.8 and up.

I spent more time on this than I wanted, because the other answers treated the whole multidimensional array as a single search where top_k only looks at the last dimension. There's more information here, where the partition is used to specifically sort a given axis.

To summarize, based upon the tensorflow signature (without name):

def top_k(input, k=1, sorted=True):
    """Top k max pooling
    Args:
        input(ndarray): convolutional feature in heigh x width x channel format
        k(int): if k==1, it is equal to normal max pooling
        sorted(bool): whether to return the array sorted by channel value
    Returns:
        ndarray: k x (height x width)
        ndarray: k
    """
    ind = np.argpartition(input, -k)[..., -k:]
    def get_entries(input, ind, sorted):
        if len(ind.shape) == 1:
            if sorted:
                ind = ind[np.argsort(-input[ind])]
            return input[ind], ind
        output, ind = zip(*[get_entries(inp, id, sorted) for inp, id in zip(input, ind)])
        return np.array(output), np.array(ind)
    return get_entries(input, ind, sorted)

Keep in mind, for your answer, you tested with

arr =  np.random.rand(3, 3, 3)
arr1, ind1 = top_k(arr)
arr2 = np.max(arr, axis=(0,1)) 
arr3, ind3 = tf.nn.top_k(arr)
print(arr1)
print(arr2)
print(arr3.numpy())

but arr2.shape is (3,) and arr3.numpy().shape is (3, 3, 1).

If you really want tf.nn.top_k like functionality, you should use np.array_equal(arr3, np.max(arr, axis=-1, keepdims=True)) as the test. I ran this with tf.enable_eager_execution() executed, hence the .numpy() instead of .eval().

Poik
  • 2,022
  • 27
  • 44
  • I think this code has a bug. What is ```a``` in your code? – neouyghur Jan 24 '19 at 02:47
  • still can't test your code. There is an error. Could you test your code as I did in my method? Thank.s – neouyghur Jan 24 '19 at 06:31
  • Your method doesn't match `tf.nn.top_k` results. There were a few problems with my code though. And I'm definitely not yet satisfied with the results, but at the very least, argpartition instead of argsort reduces the complexity by O(log(n)), so it's a start. – Poik Jan 24 '19 at 17:13
0
import numpy as np

def top_k(input, k=1):
    """Top k max pooling
    Args:
        input(ndarray): convolutional feature in heigh x width x channel format
        k(int): if k==1, it is equal to normal max pooling
    Returns:
        ndarray: k x (height x width)
    """
    input  = np.reshape(input, [-1, input.shape[-1]])
    input = np.sort(input, axis=0)[::-1, :][:k, :]
    return input


arr =  np.random.rand(3, 3, 3)
arr1 = top_k(arr)
arr2 = np.max(arr, axis=(0,1))
print(arr1)
print(arr2)
assert np.array_equal(top_k(arr)[0], np.max(arr, axis=(0,1)))
neouyghur
  • 1,577
  • 16
  • 31