如何用Numpy实现tf.nn.top_k?

时间:2018-09-02 03:55:13

标签: python numpy

如何使用Numpy实现tensorflow函数tf.nn.top_k?假设输入格式为高x宽x通道的ndarray?

2 个答案:

答案 0 :(得分: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)))

答案 1 :(得分:0)

您可以在Numpy 1.8及更高版本中使用答案here

我花了比我想要的更多的时间,因为其他答案将整个多维数组视为一次搜索,其中top_k仅查看最后一个维度。有更多信息here,其中该分区用于对给定轴进行特定排序。

总而言之,基于张量流签名(无名称):

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)

请记住,您的答案是通过

进行测试的
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())

但是arr2.shape(3,),而arr3.numpy().shape(3, 3, 1)

如果您确实想要tf.nn.top_k类似的功能,则应使用np.array_equal(arr3, np.max(arr, axis=-1, keepdims=True))作为测试。我在执行tf.enable_eager_execution()的情况下运行了此程序,因此是.numpy()而不是.eval()