train.py

#!/usr/bin/python3

# -*- coding: utf-8 -*-
"""tensorflow/datasets

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/github/tensorflow/datasets/blob/master/docs/keras_example.ipynb

# Training a neural network on MNIST with Keras

This simple example demonstrate how to plug TFDS into a Keras model.

Copyright 2020 The TensorFlow Datasets Authors, Licensed under the Apache License, Version 2.0

<table class="tfo-notebook-buttons" align="left">
  <td>
    <a target="_blank" href="https://www.tensorflow.org/datasets/keras_example"><img src="https://www.tensorflow.org/images/tf_logo_32px.png" />View on TensorFlow.org</a>
  </td>
  <td>
    <a target="_blank" href="https://colab.research.google.com/github/tensorflow/datasets/blob/master/docs/keras_example.ipynb"><img src="https://www.tensorflow.org/images/colab_logo_32px.png" />Run in Google Colab</a>
  </td>
  <td>
    <a target="_blank" href="https://github.com/tensorflow/datasets/blob/master/docs/keras_example.ipynb"><img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />View source on GitHub</a>
  </td>
</table>
"""

import tensorflow.compat.v2 as tf
import tensorflow_datasets as tfds
import sys
import re

tfds.disable_progress_bar()
tf.enable_v2_behavior()

"""## Step 1: Create your input pipeline

Build efficient input pipeline using advices from:
* [TFDS performance guide](https://www.tensorflow.org/datasets/performances)
* [tf.data performance guide](https://www.tensorflow.org/guide/data_performance#optimize_performance)

### Load MNIST

Load with the following arguments:

* `shuffle_files`: The MNIST data is only stored in a single file, but for larger datasets with multiple files on disk, it's good practice to shuffle them when training.
* `as_supervised`: Returns tuple `(img, label)` instead of dict `{'image': img, 'label': label}`
"""

dataset = 'fashion_mnist' if 'fashion' in sys.argv else 'mnist'

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

"""### Build training pipeline

Apply the following transormations:

* `ds.map`: TFDS provide the images as tf.uint8, while the model expect tf.float32, so normalize images
* `ds.cache` As the dataset fit in memory, cache before shuffling for better performance.<br/>
__Note:__ Random transformations should be applied after caching
* `ds.shuffle`: For true randomness, set the shuffle buffer to the full dataset size.<br/>
__Note:__ For bigger datasets which do not fit in memory, a standard value is 1000 if your system allows it.
* `ds.batch`: Batch after shuffling to get unique batches at each epoch.
* `ds.prefetch`: Good practice to end the pipeline by prefetching [for performances](https://www.tensorflow.org/guide/data_performance#prefetching).
"""

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.experimental.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.experimental.AUTOTUNE)

"""### Build evaluation pipeline

Testing pipeline is similar to the training pipeline, with small differences:

 * No `ds.shuffle()` call
 * Caching is done after batching (as batches can be the same between epoch)
"""

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

"""## Step 2: Create and train the model

Plug the input pipeline into Keras.
"""

network = 'A'
n_epochs = 10
lr = 0.1
adam = False
amsgrad = False

if len(sys.argv) > 1:
  network = sys.argv[1]
if len(sys.argv) > 2:
  n_epochs = int(sys.argv[2])
if len(sys.argv) > 3:
  adam = re.match('adam(.*)', sys.argv[3])
  amsgrad = re.match('amsgrad(.*)', sys.argv[3])
  if adam or amsgrad:
    try:
      lr = float((adam or amsgrad).group(1))
    except:
      lr = .001
    if adam:
      print('use Adam with lr', lr)
    else:
      print('use AMSGrad with lr', lr)
  else:
    lr = float(sys.argv[3])

if network == 'A':
  layers = [
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(128,activation='relu'),
    tf.keras.layers.Dense(128,activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
  ]
elif network == 'D':
  layers = [
    tf.keras.layers.Conv2D(5, 5, 2, 'same', activation='relu'),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(100, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
  ]
  model = tf.keras.models.Sequential(layers)
elif network == 'B':
  layers = [
    tf.keras.layers.Conv2D(16, 5, 1, 'valid', activation='relu'),
    tf.keras.layers.MaxPooling2D(2),
    tf.keras.layers.Conv2D(16, 5, 1, 'valid', activation='relu'),
    tf.keras.layers.MaxPooling2D(2),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(100, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
  ]
elif network == 'C':
  layers = [
    tf.keras.layers.Conv2D(20, 5, 1, 'valid', activation='relu'),
    tf.keras.layers.MaxPooling2D(2),
    tf.keras.layers.Conv2D(50, 5, 1, 'valid', activation='relu'),
    tf.keras.layers.MaxPooling2D(2),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(500, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
  ]
  if 'dropout' in sys.argv:
    layers.insert(-2, tf.keras.layers.Dropout(0.5))
else:
  raise Exception('unknown network: ' + network)

print (layers)
model = tf.keras.models.Sequential(layers)

if adam or amsgrad:
  optim = tf.keras.optimizers.Adam(lr, amsgrad=amsgrad)
else:
  optim = tf.keras.optimizers.SGD(momentum=0.9, learning_rate=lr,)

model.compile(
  loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
  metrics=[tf.keras.metrics.SparseCategoricalAccuracy()],
  optimizer=optim
)

for i in range(1):
  history = model.fit(
    ds_train,
    epochs=n_epochs,
    validation_data=ds_test,
  )

  out = open('log-' + '-'.join(sys.argv[1:]), 'w')
  for i in range(n_epochs):
    out.write(str(i) + ' ')
    for x in history.history:
      out.write('%s: %.4f ' % (x, history.history[x][i]))
    out.write('\n')