博客

详解TensorFlow的 tf.nn.top_k 函数：求张量中的 top-k 值和索引

  • Post category:Python

TensorFlow中的tf.nn.top_k函数

TensorFlow中的tf.nn.top_k函数可以用于计算数据张量中某个维度上前k个最大的值以及对应的索引。tf.nn.top_k具有以下方法：

tf.nn.top_k(input, k=1, sorted=True, name=None)

函数参数解析

  • input：需要进行获取top的tensor。
  • k：获取top的值数。
  • sorted：True时，按照数值排序；False时，按照原始顺序排序。
  • name：该操作的名称（可选）。

使用示例

以下是两个使用示例，以帮助更好地理解tf.nn.top_k函数：

示例1：

import tensorflow as tf

a = tf.constant([1, 3, 5, 10, 30, 20, 15])
values, indices = tf.nn.top_k(a, 3)

with tf.Session() as sess:
    print(sess.run(values))
    print(sess.run(indices))

输出结果如下：

[30 20 15]
[4 5 6]

解释：在这个示例中，我们将一个包含7个值的一维张量传入tf.nn.top_k函数，并指定k值为3。结果，我们得到了一个包含前三个最大值的一维张量（输出values），以及它们在原始输入张量中的索引位置（输出indices）。

示例2：

import tensorflow as tf

b = tf.constant([[1., 5., 3.],[7., 9., 6.],[4., 2., 8.]])
values, indices = tf.nn.top_k(b, 2)

with tf.Session() as sess:
    print(sess.run(values))
    print(sess.run(indices))

输出结果如下：

[[ 5.  3.]
 [ 9.  7.]
 [ 8.  4.]]
[[1 2]
 [1 0]
 [2 0]]

解释：在这个示例中，我们将一个包含3行3列的二维张量传入tf.nn.top_k函数，并指定k值为2。结果，我们得到了一个包含每行前两个最大值的二维张量（输出values），以及它们在原始输入张量中的索引位置（输出indices）。

可以看出，通过tf.nn.top_k函数，我们可以很快地获得输入张量中前k个最大值以及对应的索引位置。这个函数在神经网络中有很多应用，如获取预测结果中前k个最大可能性的类别。

Tags: