计算 ranking loss 时候可能需要对 tensor 进行排序和获取对应的index,
下面介绍利用 tf.nn.top_k 进行排序和获取序号
有一个 tensor,[9, 4, 5], 想得到它从小到大的排序的index [1, 2, 0]
tf.nn.top_k 可以进行从大到小排序,只要加个负号就可以进行从小到大排序了
import tensorflow as tf
a = tf.constant([3, 4, -1, 6])
size_a = tf.size(a)
max_index = tf.nn.top_k(a, size_a)[1]
min_index = tf.nn.top_k(-a, size_a)[1]
sess = tf.Session()
idx_max, idx_min = sess.run([index_max, index_min])
print(idx_max, idx_min)
显示的结果是
(array([3, 1, 0, 2], dtype=int32), array([2, 0, 1, 3], dtype=int32))
以下是 tf.nn.top_k 的API:
tf.nn.top_k(input, k=1, sorted=True, name=None)
Args:
- predictions: A Tensor of type float32. A batch_size x classes tensor.
- targets: A Tensor. Must be one of the following types: int32, int64. A batch_size vector of class ids.
- k: An int. Number of top elements to look at for computing precision.
- name: A name for the operation (optional).
另外,也可以利用 tf.reverse 对数组进行逆序
min_index = tf.reverse(max_index, axis=[0])
tf.reverse(tensor, axis, name=None)
Args:
- tensor: A Tensor. Must be one of the following types: uint8, int8, int32, int64, bool, half, float32, float64, complex64, complex128. Up to 8-D.
- axis: A Tensor. Must be one of the following types: int32, int64. 1-D. The indices of the dimensions to reverse.
- name: A name for the operation (optional).