Advanced TensorBoard Visualization for MNIST Classification with TensorFlow

The MNIST dataset consists of 28×28 grayscale images of handwritten digits; each image is flattened into a 784‑dimensional vector. The article shows an example of the digit "1" after this transformation.

A neural network is constructed with a softmax activation for the output layer because softmax has been widely proven suitable for multi‑class classification tasks. The network architecture diagram and its mathematical formulation are presented.

The model predicts a probability distribution y for the ten digit classes, while the true distribution y' is the one‑hot encoded label. The cross‑entropy loss is defined as

mean_cross_entropy = tf.reduce_mean(-tf.reduce_sum(ys * tf.log(prediction), reduction_indices=1), name='loss')

Training uses a GradientDescentOptimizer with a learning rate of 0.5:

train_op = tf.train.GradientDescentOptimizer(0.5).minimize(mean_cross_entropy)  # 0.5 is learning rate

TensorBoard visualizations are added by wrapping operations in tf.name_scope('name'), then using tf.summary.scalar('name', tensor) for scalar values such as loss, and tf.summary.histogram('name', layer) for full layer tensors (e.g., hidden and output layers). The summaries become visible only after executing sess.run(tf.summary.merge_all()).

The complete Python script (TensorFlow 1.x) includes:

# -*- coding: utf-8 -*-
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data

mnist = input_data.read_data_sets('MNIST_data', one_hot=True)
print(mnist.train.images.shape)  # (55000, 784)
print(mnist.train.labels.shape)  # (55000, 10)

def add_layer(inputs, in_size, out_size, activation_fucntion=None, name='layer'):
    with tf.name_scope(name):
        with tf.name_scope('weights'):
            Weights = tf.Variable(tf.random_normal([in_size, out_size]))
        with tf.name_scope('biases'):
            biases = tf.Variable(tf.zeros([1, out_size]) + 0.1)
        with tf.name_scope('Wx_plus_b'):
            Wx_plus_b = tf.matmul(inputs, Weights) + biases
        outputs = activation_fucntion(Wx_plus_b) if activation_fucntion else Wx_plus_b
        return outputs

def compute_accuracy(validate_xs, validate_ys):
    y_preiction = sess.run(prediction, feed_dict={xs: validate_xs})
    correct_prediction = tf.equal(tf.argmax(y_preiction, 1), tf.argmax(validate_ys, 1))
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
    return sess.run(accuracy, feed_dict={xs: validate_xs, ys: validate_ys})

with tf.name_scope('inputs'):
    xs = tf.placeholder(tf.float32, [None, 784], name='input_x')
    ys = tf.placeholder(tf.float32, [None, 10], name='input_y')

with tf.variable_scope('Net'):
    hidden1 = add_layer(xs, 784, 784*2, activation_fucntion=tf.nn.softmax, name='hidden1')
    tf.summary.histogram('hidden1', hidden1)
    prediction = add_layer(hidden1, 784*2, 10, activation_fucntion=tf.nn.softmax, name='output')
    tf.summary.histogram('prediction', prediction)

with tf.name_scope('loss'):
    mean_cross_entropy = tf.reduce_mean(-tf.reduce_sum(ys * tf.log(prediction), reduction_indices=1), name='loss')
    tf.summary.scalar('loss', mean_cross_entropy)

with tf.name_scope('train'):
    train_op = tf.train.GradientDescentOptimizer(0.5).minimize(mean_cross_entropy)

sess = tf.Session()
writer = tf.summary.FileWriter('logs/', sess.graph)
merge_op = tf.summary.merge_all()
init = tf.global_variables_initializer()
sess.run(init)

for step in range(1000):
    batch_xs, batch_ys = mnist.train.next_batch(100)
    _, result = sess.run([train_op, merge_op], feed_dict={xs: batch_xs, ys: batch_ys})
    writer.add_summary(result, step)
    if step % 50 == 0:
        print(compute_accuracy(mnist.test.images, mnist.test.labels))

After running the script, launch TensorBoard with the command: tensorboard --logdir logs TensorBoard starts at http://0.0.0.0:6006 (or http://localhost:6006) where the scalar loss curve and histograms of the hidden and output layers can be inspected.