In [1]:
# Load the TensorBoard notebook extension
%load_ext tensorboard
In [2]:
import os
import datetime
import tensorflow as tf

from tensorflow.keras import Model
from tensorflow.keras.layers import Dense, Flatten, Conv2D
In [3]:
# Load the digit dataset
mnist = tf.keras.datasets.mnist

# Spliting data into train and test sets
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

# Add a channels dimension
x_train = x_train[..., tf.newaxis]
x_test = x_test[..., tf.newaxis]
In [4]:
train_dataset  = tf.data.Dataset.from_tensor_slices(
    (x_train, y_train)).shuffle(10000).batch(32)

test_dataset  = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(32)
In [5]:
# Creating a model
class ConvNet(Model):
    def __init__(self):
        super(ConvNet, self).__init__()
        # Convolutional layer with 32 filters
        self.conv1 = Conv2D(32, 3, activation='relu')
        # Flatten - Flattens the input. Does not affect the batch size
        self.flatten = Flatten()
        # Dense - Fully conenected layer
        self.d1 = Dense(128, activation='relu')
        self.d2 = Dense(10, activation='softmax')

    def call(self, x):
        x = self.conv1(x)
        x = self.flatten(x)
        x = self.d1(x)

        return self.d2(x)
    
    # Function to trace in order to visualize the graph
    @tf.function
    def traceme(self, x):
        x = self.conv1(x)
        x = self.flatten(x)
        x = self.d1(x)

        return self.d2(x)
In [6]:
# Create an instance of the model
model = ConvNet()
In [7]:
# Hyperparameters
training_epochs = 5 # Total number of training epochs
learning_rate = 0.001 # The learning rate
In [8]:
# Set the optimizer with desired learning rate
optimizer = tf.keras.optimizers.Adam(learning_rate = learning_rate)

# We are going to use the crossentropy loss between the labels and predictions
loss_object = tf.keras.losses.SparseCategoricalCrossentropy()
In [9]:
# Select the metrics for training and test

train_loss = tf.keras.metrics.Mean(name='train_loss')
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')

test_loss = tf.keras.metrics.Mean(name='test_loss')
test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='test_accuracy')
In [10]:
# Select the metrics for training and test

@tf.function
def train_step(images, labels):
    with tf.GradientTape() as tape:
        predictions = model(images)
        loss = loss_object(labels, predictions)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    train_loss(loss)
    train_accuracy(labels, predictions)
In [11]:
# Define the test function

@tf.function
def test_step(images, labels):
    predictions = model(images)
    t_loss = loss_object(labels, predictions)

    test_loss(t_loss)
    test_accuracy(labels, predictions)
In [12]:
# Set up loggings
# In order not to overwrite some data, place the logs in a timestamped subdirectory
# This allows to easy select different runs
current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")

train_log_dir = "logs\\gradient_tape\\" + current_time + "\\train"
test_log_dir = "logs\\gradient_tape\\" + current_time + "\\test"
graph_log_dir = "logs\\gradient_tape\\" + current_time + "\\graph"

if not os.path.exists(train_log_dir):
    os.makedirs(train_log_dir)
    
if not os.path.exists(test_log_dir):
    os.makedirs(test_log_dir)
    
if not os.path.exists(graph_log_dir):
    os.makedirs(graph_log_dir)


train_summary_writer = tf.summary.create_file_writer(train_log_dir)
test_summary_writer = tf.summary.create_file_writer(test_log_dir)
graph_summary_writer = tf.summary.create_file_writer(graph_log_dir)
In [13]:
# Start the training.

for epoch in range(training_epochs):
    for images, labels in train_dataset :
        train_step(images, labels)
    with train_summary_writer.as_default():
        # Using tf.summary.scalar to log metrics during training
        tf.summary.scalar('loss', train_loss.result(), step=epoch)
        tf.summary.scalar('accuracy', train_accuracy.result(), step=epoch)

    for test_images, test_labels in test_dataset :
        test_step(test_images, test_labels)
    with test_summary_writer.as_default():
        # Using tf.summary.scalar to log metrics during testing
        tf.summary.scalar('loss', test_loss.result(), step=epoch)
        tf.summary.scalar('accuracy', test_accuracy.result(), step=epoch)

    Accuracy = 'Epoch {}, Loss: {}, Accuracy: {}, Test Loss: {}, Test Accuracy: {}'
    print(Accuracy.format(epoch+1,
                        train_loss.result(),
                        train_accuracy.result()*100,
                        test_loss.result(),
                        test_accuracy.result()*100))

    # Reset the metrics for the next epoch
    train_loss.reset_states()
    train_accuracy.reset_states()
    test_loss.reset_states()
    test_accuracy.reset_states()

WARNING:tensorflow:Layer conv_net is casting an input tensor from dtype float64 to the layer's dtype of float32, which is new behavior in TensorFlow 2.  The layer has dtype float32 because it's dtype defaults to floatx.

If you intended to run this layer in float32, you can safely ignore this warning. If in doubt, this warning is likely only an issue if you are porting a TensorFlow 1.X model to TensorFlow 2.

To change all layers to have dtype float64 by default, call `tf.keras.backend.set_floatx('float64')`. To change just this layer, pass dtype='float64' to the layer constructor. If you are the author of this layer, you can disable autocasting by passing autocast=False to the base Layer constructor.

Epoch 1, Loss: 0.14156417548656464, Accuracy: 95.76499938964844, Test Loss: 0.06890291720628738, Test Accuracy: 97.7699966430664
Epoch 2, Loss: 0.04590185731649399, Accuracy: 98.57666778564453, Test Loss: 0.05793898552656174, Test Accuracy: 98.0
Epoch 3, Loss: 0.024010632187128067, Accuracy: 99.2066650390625, Test Loss: 0.05671333521604538, Test Accuracy: 98.11000061035156
Epoch 4, Loss: 0.014908695593476295, Accuracy: 99.5183334350586, Test Loss: 0.046168532222509384, Test Accuracy: 98.68000030517578
Epoch 5, Loss: 0.009632728062570095, Accuracy: 99.66500091552734, Test Loss: 0.05531631037592888, Test Accuracy: 98.47999572753906
In [14]:
# In order to view the graph, we will trace
# how some image(tensor) flows through graph

# Enagle tracing
tf.summary.trace_on(graph=True, profiler=True)
# Forward pass
model.traceme(tf.zeros((1, 28, 28, 1)))
# Log traced path for visualization
with graph_summary_writer.as_default():
    tf.summary.trace_export(name="model_trace", step=0, profiler_outdir=graph_log_dir)

Start TensorBoard with the line magic function. For command line, run the same command without "%"

In [15]:
%tensorboard --logdir logs/gradient_tape