pdfcoffee

Chapter 16
Loading data with tf.data
Our goal is to implement a simple CNN on MNIST data (see Chapter 4,
Convolutional Neural Networks). Then we want to run the model on a TPU. To
do this, we must load the data with tf.data libraries. Hence, we need to define
a training and test function (see Chapter 2, TensorFlow 1.x and 2.x) as shown in the
following code:
# training input function
def train_input_fn(batch_size=1024):
# Convert the inputs to a Dataset.
dataset = tf.data.Dataset.from_tensor_slices((x_train,y_train))
# Shuffle, repeat, and batch the examples.
dataset = dataset.cache() # Loads the data into memory
dataset = dataset.shuffle(1000, reshuffle_each_iteration=True)
dataset = dataset.repeat()
dataset = dataset.batch(batch_size, drop_remainder=True)
return dataset
# testing input function
def test_input_fn(batch_size=1024):
dataset = tf.data.Dataset.from_tensor_slices((x_test,y_test))
# Shuffle, repeat, and batch the examples.
dataset = dataset.cache()
dataset = dataset.shuffle(1000, reshuffle_each_iteration=True)
dataset = dataset.repeat()
dataset = dataset.batch(batch_size, drop_remainder=True)
return dataset
Where (x_train, y_train), (x_test, y_test) = mnist.load_data(). Note
that drop_remainder=True is an important parameter that forces the batch method
to pass fixed shapes expected by the TPUs. Note that TPU v2 has an MMU with 128
× 128 multipliers. Usually you get the best performance by setting the batch size to
128 per TPU core. With 10 TPU cores, for instance, the batch size would be 1,280.
Building a model and loading it into the TPU
As of November 2019, TensorFlow 2.0 does not fully support TPUs. They are
available in TensorFlow 1.5.0, and with the nightly build. Let's first see an example
with TensorFlow 1.5, and the example with the nightly build will be shown later.
[ 581 ]