mnist_train.py

import tensorflow as tf
import pandas as pd
import os
print("TensorFlow version:", tf.__version__)

mnist = tf.keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0


model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  tf.keras.layers.Dense(128, activation='relu'),
  tf.keras.layers.Dropout(0.2),
  tf.keras.layers.Dense(10)
])

loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

model.compile(optimizer='adam',
              loss=loss_fn,
              metrics=['accuracy'])

history = model.fit(x_train, y_train, epochs=15)


############################################################################
# Save a model

model.save(os.path.join("../models/","v1_mnist.h5"))


#########################################################
# Save history of training


df = pd.DataFrame(history.history)
filename = '../history/v1_mnist.csv'
with open (filename, mode ='w') as f:
    df.to_csv(f)