# Train a model
import tensorflow as tf
tf.executing_eagerly()

(mnist_images, mnist_labels), _ = tf.keras.datasets.mnist.load_data()

dataset = tf.data.Dataset.from_tensor_slices(
    (tf.cast(mnist_images[..., tf.newaxis] / 255, tf.float32),
     tf.cast(mnist_labels, tf.int64))
)
dataset = dataset.shuffle(1000).batch(32)

mnist_model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(16, [3, 3], activation='relu', input_shape=(None, None, 1)),
    tf.keras.layers.Conv2D(16, [3, 3], activation='relu'),
    tf.keras.layers.GlobalAveragePooling2D(),
    tf.keras.layers.Dense(10)
])


for images, labels in dataset.take(1):
    print("Logits: ", mnist_model(images[0:1]).numpy())

optimizer = tf.keras.optimizers.Adam()
loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

loss_history = []


def train_step(images, labels):
    with tf.GradientTape() as tape:
        logits = mnist_model(images, training=True)
        
        # Add asserts to check the shape of the output.
        tf.debugging.assert_equal(logits.shape, (32, 10))
        loss_value = loss_object(labels, logits)

    loss_history.append(loss_value.numpy().mean())
    grads = tape.gradient(loss_value, mnist_model.trainable_variables)
    optimizer.apply_gradients(zip(grads, mnist_model.trainable_variables))


def train():
    for epoch in range(3):
        for (batch, (images, labels)) in enumerate(dataset):
            train_step(images, labels)
        print('Epoch {} finished'.format(epoch))

train()

import matplotlib.pyplot as plt
plt.plot(loss_history)
plt.xlabel('Batch #')
plt.ylabel('Loss [entropy]')
plt.show()
        

tensorflow.debugging.assert_equal tensorflow.executing_eagerly tensorflow.keras.layers.GlobalAveragePooling2D tensorflow.keras.losses.SparseCategoricalCrossentropy tensorflow.keras.layers.Dense tensorflow.cast tensorflow.keras.layers.Conv2D tensorflow.keras.datasets.mnist.load_data matplotlib.pyplot.plot tensorflow.GradientTape tensorflow.keras.optimizers.Adam matplotlib.pyplot.xlabel matplotlib.pyplot.show matplotlib.pyplot.ylabel
[(3, 'tensorflow.executing_eagerly', 'tf.executing_eagerly', ([], {}), True, 'import tensorflow as tf\n'), (5, 'tensorflow.keras.datasets.mnist.load_data', 'tf.keras.datasets.mnist.load_data', ([], {}), True, 'import tensorflow as tf\n'), (24, 'tensorflow.keras.optimizers.Adam', 'tf.keras.optimizers.Adam', ([], {}), True, 'import tensorflow as tf\n'), (25, 'tensorflow.keras.losses.SparseCategoricalCrossentropy', 'tf.keras.losses.SparseCategoricalCrossentropy', ([], {'from_logits': '(True)'}), True, 'import tensorflow as tf\n'), (52, 'matplotlib.pyplot.plot', 'plt.plot', (['loss_history'], {}), True, 'import matplotlib.pyplot as plt\n'), (53, 'matplotlib.pyplot.xlabel', 'plt.xlabel', (['"""Batch #"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (54, 'matplotlib.pyplot.ylabel', 'plt.ylabel', (['"""Loss [entropy]"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (55, 'matplotlib.pyplot.show', 'plt.show', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (8, 'tensorflow.cast', 'tf.cast', (['(mnist_images[..., tf.newaxis] / 255)', 'tf.float32'], {}), True, 'import tensorflow as tf\n'), (9, 'tensorflow.cast', 'tf.cast', (['mnist_labels', 'tf.int64'], {}), True, 'import tensorflow as tf\n'), (14, 'tensorflow.keras.layers.Conv2D', 'tf.keras.layers.Conv2D', (['(16)', '[3, 3]'], {'activation': '"""relu"""', 'input_shape': '(None, None, 1)'}), True, 'import tensorflow as tf\n'), (15, 'tensorflow.keras.layers.Conv2D', 'tf.keras.layers.Conv2D', (['(16)', '[3, 3]'], {'activation': '"""relu"""'}), True, 'import tensorflow as tf\n'), (16, 'tensorflow.keras.layers.GlobalAveragePooling2D', 'tf.keras.layers.GlobalAveragePooling2D', ([], {}), True, 'import tensorflow as tf\n'), (17, 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', (['(10)'], {}), True, 'import tensorflow as tf\n'), (31, 'tensorflow.GradientTape', 'tf.GradientTape', ([], {}), True, 'import tensorflow as tf\n'), (35, 'tensorflow.debugging.assert_equal', 'tf.debugging.assert_equal', (['logits.shape', '(32, 10)'], {}), True, 'import tensorflow as tf\n')]