使用argmax从张量获取值

问题描述

我有一个Tensor形状(60, 128, 30000)。我想获取30000维度(axis=2)的argmax的值。 以下代码是一个示例：

tensor = tf.random.uniform((60, 128, 30000)) # shape (60, 128, 30000)
argmax = tf.argmax(tensor, axis=2) # shape (60, 128) --> max of each 30000

# do something to get every values of 30000

# argmax output (index)
<tf.Tensor: shape=(60, 128), dtype=int64, numpy=
array([[ 3229,  3079,  8360, ...,  1005, 16460,   872],
       [17808,  1253, 25476, ..., 16130,  3479,  3479],
       [27717, 25429, 18808, ...,  9787,  2603, 24011],
       ...,
       [25429, 25429,  5647, ..., 18451, 12453, 12453],
       [ 7361, 13463, 15864, ..., 18839, 12453, 12453],
       [ 4750, 25009, 11888, ...,  5647,  1993, 18451]], dtype=int64)>

# Desired output: each values of every index

使用argmax，我得到的是它们的索引数组，而不是它们的值。如何获取它们的值的形状相同的数组(60, 128)？

---

解决方案

您必须使用tf.meshgrid和tf.gather_nd来实现您想要的效果：

tensor = tf.random.uniform((60, 128, 30000)) # shape (60, 128, 30000)
argmax = tf.argmax(tensor, axis=2)

ij = tf.stack(tf.meshgrid(
    tf.range(tensor.shape[0], dtype=tf.int64), 
    tf.range(tensor.shape[1], dtype=tf.int64),
                              indexing='ij'), axis=-1)

gather_indices = tf.concat([ij, tf.expand_dims(argmax, axis=-1)], axis=-1)
result = tf.gather_nd(tensor, gather_indices)
tf.print(result.shape)

TensorShape([60, 128])

为什么需要tf.meshgrid？因为argmax确实包含您的索引，但形式错误。函数tf.gather_nd需要知道它应该准确地从3D张量中提取值的位置。tf.meshgrid函数用于创建由两个一维数组组成的矩形网格，表示第一维和第二维的张量索引。

import tensorflow as tf

tensor = tf.random.uniform((2, 5, 3))
argmax = tf.argmax(tensor, axis=2)

# result = tf.gather_nd(tensor, gather_ind) <-- Would not work because arxmax has the shape TensorShape([2, 5]) but  TensorShape([2, 5, 3]) is required
tf.print('Input tensor:
', tensor, tensor.shape, '
Argmax tensor:
', argmax, argmax.shape)

i, j = tf.meshgrid(
    tf.range(tensor.shape[0], dtype=tf.int64), 
    tf.range(tensor.shape[1], dtype=tf.int64),
                              indexing='ij')

# You need to create a mesh grid to correctly index your tensor.

ij = tf.stack([i, j], axis=-1)
tf.print('Meshgrid:
', i, j, summarize=-1)
tf.print('Stacked:
', ij, summarize=-1)

gather_indices = tf.concat([ij, tf.expand_dims(argmax, axis=-1)], axis=-1)
tf.print('Gathered indices:
', gather_indices, gather_indices.shape, summarize=-1)

result = tf.gather_nd(tensor, gather_indices)
tf.print('
Final result:
', result, result.shape)

Input tensor:
 [[[0.889752269 0.243187189 0.601408958]
  [0.891950965 0.776625633 0.146243811]
  [0.136176467 0.743871331 0.762170076]
  [0.424416184 0.150568008 0.464055896]
  [0.308753 0.0792338848 0.383242]]

 [[0.741660118 0.49783361 0.935318112]
  [0.0616152287 0.0367363691 0.748341084]
  [0.397849679 0.765681744 0.502376914]
  [0.750188231 0.304993749 0.733741879]
  [0.31267941 0.778184056 0.546301]]] TensorShape([2, 5, 3]) 
Argmax tensor:
 [[0 0 2 2 2]
 [2 2 1 0 1]] TensorShape([2, 5])
Meshgrid:
 [[0 0 0 0 0]
 [1 1 1 1 1]] [[0 1 2 3 4]
 [0 1 2 3 4]]
Stacked:
 [[[0 0]
  [0 1]
  [0 2]
  [0 3]
  [0 4]]

 [[1 0]
  [1 1]
  [1 2]
  [1 3]
  [1 4]]]
Gathered indices:
 [[[0 0 0]
  [0 1 0]
  [0 2 2]
  [0 3 2]
  [0 4 2]]

 [[1 0 2]
  [1 1 2]
  [1 2 1]
  [1 3 0]
  [1 4 1]]] TensorShape([2, 5, 3])

Final result:
 [[0.889752269 0.891950965 0.762170076 0.464055896 0.383242]
 [0.935318112 0.748341084 0.765681744 0.750188231 0.778184056]] TensorShape([2, 5])

顺便说一句，您还可以考虑使用tf.math.top_k，因为您希望获得最后一个维度的最大值。此函数返回您想要的索引和值：

tensor = tf.random.uniform((60, 128, 30000)) # shape (60, 128, 30000)
values, indices = tf.math.top_k(tensor,
                        k=1)
tf.print(tf.squeeze(values, axis=-1).shape)

TensorShape([60, 128])