如何在tf.Tensor中按指定维度提取每个子元素的最大值

问题原因 #

你定义的张量a形状为(4, 1, 2)，调用tf.reduce_max时未指定运算轴参数axis，函数默认会对张量的所有维度做全局最大值聚合，因此只会返回整个张量的全局最大值8。

解决方案 #

你需要指定对存储子元素数值的最后一个维度（索引为-1或者2）做最大值聚合：

import tensorflow as tf
a = tf.constant([
  [[1, 2]], [[3, 4]],
  [[5, 6]], [[7, 8]]])

# 得到形状为(4,1)的结果，每个位置对应子元素的最大值
res = tf.reduce_max(a, axis=-1)
print(res)
# 输出：
# <tf.Tensor: shape=(4, 1), dtype=int32, numpy=
# array([[2],
#        [4],
#        [6],
#        [8]])>

如果需要直接得到一维的[2 4 6 8]结果，可以额外调用tf.squeeze去掉多余的空维度，或者直接指定聚合轴为[1,2]：

# 方式1：先求最大值再压缩维度
res = tf.squeeze(tf.reduce_max(a, axis=-1))
# 方式2：直接对第1、2位维度做聚合
res = tf.reduce_max(a, axis=[1,2])
print(res)
# 输出：tf.Tensor([2 4 6 8], shape=(4,), dtype=int32)

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