如何在TensorFlow中获取张量最大值的索引？

给定形状为(2, 2, 3)的目标张量，要得到每个最内层子张量最大值的完整三维索引，最简单的实现方式可以结合tf.argmax和tf.meshgrid完成：

步骤分解 #

• 提取最内层维度的最大值索引：用tf.argmax指定axis=2，得到每个位置在第三维度的最大值索引，结果形状为(2,2)。
• 生成前两个维度的网格索引：通过tf.meshgrid生成对应第一、第二维度的索引矩阵，注意指定indexing='ij'来匹配原张量的维度顺序。
• 拼接并调整形状：将三个维度的索引拼接后，展平为每行对应一个完整三维索引的形式。

完整代码示例 #

import tensorflow as tf

# 定义示例张量
tensor = tf.constant([
    [[0.1, 0.3, 0.6],
     [0.0, 0.4, 0.1]],
    [[0.9, 0.2, 0.6],
     [0.8, 0.1, 0.5]]
])

# 获取第三维度的最大值索引
max_indices = tf.argmax(tensor, axis=2)

# 生成前两个维度的网格索引
dim0, dim1 = tf.meshgrid(tf.range(tensor.shape[0]), tf.range(tensor.shape[1]), indexing='ij')

# 拼接所有维度索引并调整为目标形状
full_indices = tf.stack([dim0, dim1, max_indices], axis=-1)
full_indices = tf.reshape(full_indices, (-1, 3))

# 输出结果
print(full_indices.numpy())

运行代码后会输出：

[[0 0 2]
 [0 1 1]
 [1 0 0]
 [1 1 0]]

关键说明 #

• tf.meshgrid的indexing='ij'参数是核心：它保证生成的索引矩阵和原张量的维度顺序一致（第一维度行优先），避免出现索引错位。
• tf.stack将三个维度的索引在最后一维合并，再通过tf.reshape展平，正好匹配你需要的每行一个完整索引的格式。

内容的提问来源于stack exchange，提问作者estrella1995