Tensorflow中如何通过已生成的arg_max获取张量对应的最大值

实现方案 #

方法1：基于NumPy索引实现 #

先将TensorFlow格式的arg_max转为NumPy数组，再结合行索引做高级取值，最后调整维度即可：

import numpy as np
import tensorflow as tf

# 你的原始代码
a = np.array([[0, 10], [1.1, 0]])
arg_max = tf.math.argmax(input = a, axis=1)

# 转换张量为NumPy数组
arg_max_np = arg_max.numpy()
# 按行索引+最大值列索引取值
max_values = a[np.arange(a.shape[0]), arg_max_np]
# 调整维度为(N,1)格式
result = max_values.reshape(-1, 1)

运行后result输出即为你需要的[[10. ], [ 1.1]]。

方法2：纯TensorFlow实现（无需转NumPy） #

如果全程使用TensorFlow接口操作，可以用tf.gather_nd直接取值：

a_tf = tf.convert_to_tensor(a)
# 构造坐标索引：格式为[(行号, 最大值列号)]
indices = tf.stack([tf.range(tf.shape(a_tf)[0]), arg_max], axis=1)
result_tf = tf.gather_nd(a_tf, indices)
# 调整维度为目标格式
result = tf.expand_dims(result_tf, axis=1).numpy()

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