How to Use tf.reduce_mean in TensorFlow

tf.reduce_mean

Computes the mean of elements across dimensions of a tensor.

To use this function, you must notice the axis.

import tensorflow as tf
# a: 2*3
a = tf.constant([[1, 2, 3], [4, 5, 6]], dtype=tf.float32)
b = tf.reduce_mean(a,0) # (3,)
c = tf.reduce_mean(a,1) # (2,)

# d:(2,2,3)
d = tf.constant([[[1, 2, 3], [4, 5, 6]],[[1, 6, 3], [4, 3, 6]]], dtype=tf.float32)
e = tf.reduce_mean(d,0) # (2,3)
f = tf.reduce_mean(d,1) # (2,3)
h = tf.reduce_mean(d,2) # (2,2)

with tf.Session() as sess:
    sess.run(init) 
    print(sess.run(b))
    print(sess.run(c))
    print(sess.run(e))
    print(sess.run(f))
    print(sess.run(h))

The reust is:

[2.5 3.5 4.5]
[2. 5.]
[[1. 4. 3.]
 [4. 4. 6.]]
[[2.5 3.5 4.5]
 [2.5 4.5 4.5]]
[[2.        5.       ]
 [3.3333333 4.3333335]]

,