Skip to content

Instantly share code, notes, and snippets.

@ttamg

ttamg/tensorflow_mnist.py

Last active March 13, 2022 17:06
Show Gist options
  • Save ttamg/a1f3c347b829ef9715e65c24d49eb26c to your computer and use it in GitHub Desktop.
Save ttamg/a1f3c347b829ef9715e65c24d49eb26c to your computer and use it in GitHub Desktop.
Download ZIP

Revisions

  1. ttamg revised this gist Mar 13, 2022. 1 changed file with 2 additions and 1 deletion.
    3 changes: 2 additions & 1 deletion tensorflow_mnist.py
    View file
    Original file line number Diff line number Diff line change
    @@ -1,4 +1,5 @@
    """ A simple Tensorflow 2 MNIST training script to test installation is working correctly. """
    # tensorflow_mnist.py
    # A simple Tensorflow 2 MNIST training script to test installation is working correctly.

    import tensorflow as tf
    import tensorflow_datasets as tfds
  2. ttamg created this gist Mar 13, 2022.
    46 changes: 46 additions & 0 deletions tensorflow_mnist.py
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,46 @@
    """ A simple Tensorflow 2 MNIST training script to test installation is working correctly. """

    import tensorflow as tf
    import tensorflow_datasets as tfds

    (ds_train, ds_test), ds_info = tfds.load(
    'mnist',
    split=['train', 'test'],
    shuffle_files=True,
    as_supervised=True,
    with_info=True,
    )

    def normalize_img(image, label):
    """Normalizes images: `uint8` -> `float32`."""
    return tf.cast(image, tf.float32) / 255., label

    ds_train = ds_train.map(
    normalize_img, num_parallel_calls=tf.data.AUTOTUNE)
    ds_train = ds_train.cache()
    ds_train = ds_train.shuffle(ds_info.splits['train'].num_examples)
    ds_train = ds_train.batch(128)
    ds_train = ds_train.prefetch(tf.data.AUTOTUNE)

    ds_test = ds_test.map(
    normalize_img, num_parallel_calls=tf.data.AUTOTUNE)
    ds_test = ds_test.batch(128)
    ds_test = ds_test.cache()
    ds_test = ds_test.prefetch(tf.data.AUTOTUNE)

    model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10)
    ])
    model.compile(
    optimizer=tf.keras.optimizers.Adam(0.001),
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=[tf.keras.metrics.SparseCategoricalAccuracy()],
    )

    model.fit(
    ds_train,
    epochs=6,
    validation_data=ds_test,
    )