tf.reduce_mean()的用法

tf.math.reduce_mean(
    input_tensor, axis=None, keepdims=False, name=None
)

作用：按给定维度实现平均，将指定维度大小变为1
参数：input_tensor输入张量，axis变换的维度（求和后平均）(默认值为所有元素求和取平均)，keepdims输出结果是否保持维度（默认为不保持原维度）
代码示例：
import tensorflow as tf

x = tf.constant([[[1., 1., 1], [2., 2., 2]], [[3., 3., 3], [4., 4., 4]]])
a = tf.reduce_mean(x)

b = tf.reduce_mean(x, 0)

c = tf.reduce_mean(x, 1)
d = tf.reduce_mean(x, 2)

e = tf.reduce_mean(x, [1, 2], keepdims=True)
f = tf.reduce_mean(x, [1, 2])
print(a, b, c, d, e, f)

输出结果：

tf.Tensor(2.5, shape=(), dtype=float32)

tf.Tensor(
[[2. 2. 2.]
 [3. 3. 3.]], shape=(2, 3), dtype=float32)

tf.Tensor(
[[1.5 1.5 1.5]
 [3.5 3.5 3.5]], shape=(2, 3), dtype=float32)

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

tf.Tensor(
[[[1.5]]

 [[3.5]]], shape=(2, 1, 1), dtype=float32)

tf.Tensor([1.5 3.5], shape=(2,), dtype=float32)

个人理解：所给示例为2*2*3的张量，分别对应维度0、1、2；

当axis=0时，求和后平均0所对应维度实现将该维度参数2变为1，对应示例即[[1., 1., 1], [2., 2., 2]]和[[3., 3., 3], [4., 4., 4]]对应位置求和取平均并输出，输出结果保留维度为2*3，输出结果为[[2. 2. 2.]  [3. 3. 3.]]

当axis=2时，将2*2*3的张量中维度为2对应的维度进行变换（消减该维度，3变为1），2对应维的数据分别为[1., 1., 1]...[4., 4., 4],将该维数据维度内求和平均后得到对应结果1，2，3，4，输出结果保留维度为2*2（其实就是将2对应维度剔除了），输出结果[[1. 2.]  [3. 4.]]

当axis=[1, 2]时，即对维度1和2对应进行求和平均实现维度消减，最后输出维度只保留0维度对应的大小，若keepdims=True,则输出数据保持变换后该有的维度即0对应的维度的参数大小不变（为2），1和2维度对应的参数2*3变为1*1，输出维度结果为2*1*1

个人理解