减少 tensorflow 中的最大值以对行进行分组？

Question

我有一个张量，我想在按某些行分组的每一列中计算最大值。

例如：tensor_eg =

[[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]]

tf.reduce_max(tensor_eg, axis = 0) 会给我所有行中每一列的最大值。 我想按某些行分组，例如第 0 行和第 1 行的最大值，第 2 和第 3 行的最大值：

想要的结果：

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

我怎样才能做到这一点？

Answer 1

也许只是尝试切片并使用tf.concat ：

import tensorflow as tf

x = 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]])

tf.concat([tf.reduce_max(x[:2, :], keepdims=True, axis=0), tf.reduce_max(x[2:, :], keepdims=True, axis = 0)], axis=0) 

<tf.Tensor: shape=(2, 4), dtype=float32, numpy=
array([[0.1, 0.8, 0.4, 0.5],
       [0.1, 0.2, 0.6, 0.5]], dtype=float32)>

更通用的方法是使用tf.math.segment_max ：

x = 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]])

tf.math.segment_max(x, tf.constant([0, 0, 1, 1]))

段 id 需要等于x的第一个维度的大小