如何在TensorFlow中按分组行计算列的最大值？

针对你需要按指定行分组计算每列最大值的需求，这里提供两种实用的实现方法：

方法一：连续等规模分组（如每N行一组） #

如果分组是连续且每组行数固定（比如示例中每2行一组），可以通过调整张量形状后直接聚合：

import tensorflow as tf

# 示例张量
tensor_eg = tf.constant([
    [0.1, 0.2, 0.4, 0.5],
    [0.1, 0.8, 0.2, 0.5],
    [0.1, 0.2, 0.4, 0.5],
    [0.1, 0.1, 0.6, 0.5]
])

# 将张量重塑为(组数, 每组行数, 列数)的三维结构
grouped_tensor = tf.reshape(tensor_eg, (2, 2, 4))
# 沿每组内的行维度（axis=1）计算每列最大值
result = tf.reduce_max(grouped_tensor, axis=1)

print(result.numpy())

输出结果：

[[0.1 0.8 0.4 0.5]
 [0.1 0.2 0.6 0.5]]

方法二：自定义任意分组（支持非连续/不等规模分组） #

如果分组规则更灵活（比如需要跨行分组、每组行数不同），可以使用tf.math.unsorted_segment_max函数，它能根据自定义的分组索引完成聚合计算：

import tensorflow as tf

tensor_eg = tf.constant([
    [0.1, 0.2, 0.4, 0.5],
    [0.1, 0.8, 0.2, 0.5],
    [0.1, 0.2, 0.4, 0.5],
    [0.1, 0.1, 0.6, 0.5]
])

# 定义分组索引：第0、1行归为组0，第2、3行归为组1
group_indices = tf.constant([0, 0, 1, 1])
# 计算分组最大值：参数依次为原张量、分组索引、总组数
result = tf.math.unsorted_segment_max(tensor_eg, group_indices, num_segments=2)

print(result.numpy())

同样会输出你期望的结果。这个方法的优势在于，分组索引可以任意指定，比如group_indices = [0,1,0,1]就能实现第0、2行一组，第1、3行一组的计算。

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