Understand tf.reduce_mean with Examples for Beginners – TensorFlow Tutorial

tf.reduce_mean() can allow us to compute the mean value of a tensor in tensorflow. This function is widely used in tensorflow applications. However,  to use this function correctly, we must concern how this function compute the mean of a tensor and how about the result.

Key 1. tf.reduce_mean computes the average of a tensor along axis.

(1) If we do not set axis value, tf.reduce_mean() will compute the mean of all elements in tensor.

For example:

We create a tensor with shape [5, 10, 9, 8], then compute the mean without axis.

import tensorflow as tf
import numpy as np

x = tf.Variable(tf.random_uniform([5, 10, 9, 8], -1, 1), name='x')

avg_x_4 = tf.reduce_mean(x)

init = tf.global_variables_initializer() 
init_local = tf.local_variables_initializer()
with tf.Session() as sess:
    sess.run([init, init_local])
    x0= (sess.run([avg_x_4]))
    print x0

Will get the value:

-0.008919376

(2) if we set the value of axis, tf.reduce_mean() will compute along the axis.

For example:

We set axis = 0, we will compute the mean like this:

import tensorflow as tf
import numpy as np

x = tf.Variable(tf.random_uniform([5, 10, 9, 8], -1, 1), name='x')

avg_x_4 = tf.reduce_mean(x, axis = 0)

init = tf.global_variables_initializer() 
init_local = tf.local_variables_initializer()
with tf.Session() as sess:
    sess.run([init, init_local])
    x0 = (sess.run([avg_x_4]))
    print x.shape

The value of shape is:

(10, 9, 8)

Here we will find an interesting thing, the shape of tensor x is reduced.

Key 2. If we do not set keepdims=true, the shape of x tensor will reduce one dimension, which is the axis you set.

For examle:

We create a [5, 10, 9, 8] tensor and compute its mean along aix = 0, 1, 2, 3, then we print each shape.

import tensorflow as tf
import numpy as np

x = tf.Variable(tf.random_uniform([5, 10, 9, 8], -1, 1), name='x')

avg_x_0 = tf.reduce_mean(x, 0)
avg_x_1 = tf.reduce_mean(x, 1)
avg_x_2 = tf.reduce_mean(x, 2)
avg_x_3 = tf.reduce_mean(x, 3)

init = tf.global_variables_initializer() 
init_local = tf.local_variables_initializer()
with tf.Session() as sess:
    sess.run([init, init_local])
    x, x0, x1, x2, x3= (sess.run([x, avg_x_0, avg_x_1, avg_x_2, avg_x_3]))
    
    print 'x tensor shape = '
    print x.shape
    print 'x tensor mean along axis= 0 '
    print x0.shape
    print 'x tensor mean along axis= 1 '
    print x1.shape
    print 'x tensor mean along axis= 2 '
    print x2.shape
    print 'x tensor mean along axis= 3 '
    print x3.shape

The output is:

x tensor shape = 
(5, 10, 9, 8)
x tensor mean along axis= 0 
(10, 9, 8)
x tensor mean along axis= 1 
(5, 9, 8)
x tensor mean along axis= 2 
(5, 10, 8)
x tensor mean along axis= 3 
(5, 10, 9)

We will find:

If we compute the mean of x tensor along axis = 0, the first dimension of x (5) will be removed.

Similarly, if we compute along axis = 1, 2, 3. the dimension of x (10, 9, 8) will be removed.

So

axis = 0, [5, 10, 9, 8]  will be [10, 9, 8]

axis = 1, [5, 10, 9, 8]  will be [5, 9, 8]

axis = 2, [5, 10, 9, 8]  will be [5, 10, 8]

axis = 3, [5, 10, 9, 8]  will be [5, 10, 9]