TensorFlow：如何获取批量中每个张量的单个最大值坐标

我有一个形状为[batch, height, width]的3阶张量，需要为每个batch元素（即每张图像）获取一个最大值像素的坐标，只要是该样本中的最大值坐标即可，无需保留所有重复的最大值位置。

现有实现的问题 #

以下是基于Stack Overflow回答的初始实现（假设最大值为100，且每个样本至少包含一个最大值）：

x = tf.constant([[[1, 2, 3], 
                 [4, 100, 5],
                 [6, 100, 8]],
                 
                 [[3, 100, 1], 
                 [5, 6, 7],
                 [9, 8, 7]],
                 
                 [[1, 3, 7], 
                 [2, 4, 7],
                 [100, 5, 7]],
                 
                 [[1, 3, 7], 
                 [2, 4, 7],
                 [4, 5, 100]]])


max_val = 100
cond = tf.equal(x, max_val)
res = tf.where(cond)

该代码返回的结果会包含同一batch内所有最大值的坐标（比如第0个batch有两个最大值，结果里会出现两行）：

tf.Tensor(
[[0 1 1]   
 [0 2 1]
 [1 0 1]
 [2 2 0]
 [3 2 2]], shape=(5, 3), dtype=int64)

我们需要每个batch仅保留一行坐标。

解决方案 #

方案1：对现有结果去重

如果必须基于现有tf.where的结果处理，可以通过分组取第一个匹配项实现去重：

# 提取所有结果的batch索引
batch_indices = res[:, 0]
# 获取唯一的batch标识和对应分组索引
unique_batches, group_idx = tf.unique(batch_indices)
# 找到每个分组中第一个出现的结果的索引
first_idx = tf.math.unsorted_segment_min(tf.range(tf.shape(res)[0]), group_idx, tf.shape(unique_batches)[0])
# 筛选出每个batch的第一个最大值坐标
unique_res = tf.gather(res, first_idx)

输出结果将是每个batch一行：

tf.Tensor(
[[0 1 1]
 [1 0 1]
 [2 2 0]
 [3 2 2]], shape=(4, 3), dtype=int64)

方案2：直接计算（更优，无需预设最大值）

不需要硬编码最大值，直接通过张量运算获取每个样本的最大值坐标，效率更高：

# 将每个样本的二维图像展平为一维，找到最大值的一维索引
flattened_max_idx = tf.argmax(tf.reshape(x, [tf.shape(x)[0], -1]), axis=1)
# 将一维索引转换为(height, width)二维坐标
height_idx = flattened_max_idx // tf.shape(x)[2]
width_idx = flattened_max_idx % tf.shape(x)[2]
# 组合batch索引、height和width坐标，得到最终结果
unique_res = tf.stack([tf.range(tf.shape(x)[0]), height_idx, width_idx], axis=1)

该方法自动为每个样本返回一个最大值坐标，避免了重复结果，同时不需要提前知道最大值的具体数值，适用性更强。

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