Licensed under the Apache License, Version 2.0 (the "License");


In [1]:
# Copyright 2018 The TensorFlow Hub Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

A simple classification model using Keras with Cloud TPUs

Overview

This notebook shows how to use Keras to build a simple classification model. The model can train, evaluate, and generate predictions using Cloud TPUs. It uses the iris dataset to predict the species of the flower and also shows how to use your own data instead of using pre-loaded data. This model uses 4 input features (SepalLength, SepalWidth, PetalLength, PetalWidth) to determine one of these flower species (Setosa, Versicolor, Virginica).

The model trains for 50 epochs and completes in approximately 2 minutes.

This notebook is hosted on GitHub. To view it in its original repository, after opening the notebook, select File > View on GitHub.

NOTE: This tutorial is designed to show how to write a simple model using Keras. It should not be used for comparision with training on CPU's because of the very small amount of data being used.

Learning objectives

In this Colab, you will learn how to:

  • Define a Keras model with 2 hidden layers and 10 nodes in each layer.
  • Create and compile a Keras model on TPU with a distribution strategy.
  • Train, evaluate, and and generate predictions on Cloud TPU.

Instructions

  Train on TPU  

  1. On the main menu, click Runtime and select Change runtime type. Set "TPU" as the hardware accelerator.
  2. Click Runtime again and select Runtime > Run All. You can also run the cells manually with Shift-ENTER.

Data, model, and training

Imports


In [2]:
import json
import os
import pandas as pd
import pprint
import tensorflow as tf
import time
import numpy as np
from tensorflow import keras

In [3]:
print(tf.__version__)
import distutils
if distutils.version.LooseVersion(tf.__version__) < '1.14':
    raise Exception('This notebook is compatible with TensorFlow 1.14 or higher, for TensorFlow 1.13 or lower please use the previous version at https://github.com/tensorflow/tpu/blob/r1.13/tools/colab/classification_iris_data_with_keras.ipynb')


2.2.0

Resolve TPU Address


In [4]:
use_tpu = True #@param {type:"boolean"}

if use_tpu:
    assert 'COLAB_TPU_ADDR' in os.environ, 'Missing TPU; did you request a TPU in Notebook Settings?'

if 'COLAB_TPU_ADDR' in os.environ:
  TF_MASTER = 'grpc://{}'.format(os.environ['COLAB_TPU_ADDR'])
else:
  TF_MASTER=''

FLAGS used as model params


In [5]:
# Model specific parameters

# TPU address
tpu_address = TF_MASTER

# Number of epochs
epochs = 50

# Number of steps_per_epoch
steps_per_epoch = 5

# NOTE: Total number of training steps = Number of epochs * Number of steps_per_epochs

Download training input data and define prediction input & output


In [6]:
TRAIN_URL = "http://download.tensorflow.org/data/iris_training.csv"
TEST_URL = "http://download.tensorflow.org/data/iris_test.csv"

CSV_COLUMN_NAMES = ['SepalLength', 'SepalWidth',
                    'PetalLength', 'PetalWidth', 'Species']
SPECIES = ['Setosa', 'Versicolor', 'Virginica']

PREDICTION_INPUT_DATA = {
    'SepalLength': [6.9, 5.1, 5.9, 6.0, 5.5, 6.2, 5.5, 6.3],
    'SepalWidth': [3.1, 3.3, 3.0, 3.4, 2.5, 2.9, 4.2, 2.8],
    'PetalLength': [5.4, 1.7, 4.2, 4.5, 4.0, 4.3, 1.4, 5.1],
    'PetalWidth': [2.1, 0.5, 1.5, 1.6, 1.3, 1.3, 0.2, 1.5],
}

PREDICTION_OUTPUT_DATA = ['Virginica', 'Setosa', 'Versicolor', 'Versicolor', 'Versicolor', 'Versicolor', 'Setosa', 'Virginica']

def maybe_download():
    train_path = tf.keras.utils.get_file(TRAIN_URL.split('/')[-1], TRAIN_URL)
    test_path = tf.keras.utils.get_file(TEST_URL.split('/')[-1], TEST_URL)

    return train_path, test_path

def load_data(y_name='Species'):
    """Returns the iris dataset as (train_x, train_y), (test_x, test_y)."""
    train_path, test_path = maybe_download()

    train = pd.read_csv(train_path, names=CSV_COLUMN_NAMES, header=0, dtype={'SepalLength': pd.np.float32,
        'SepalWidth': pd.np.float32, 'PetalLength': pd.np.float32, 'PetalWidth': pd.np.float32, 'Species': pd.np.int32})
    train_x, train_y = train, train.pop(y_name)

    test = pd.read_csv(test_path, names=CSV_COLUMN_NAMES, header=0, dtype={'SepalLength': pd.np.float32,
        'SepalWidth': pd.np.float32, 'PetalLength': pd.np.float32, 'PetalWidth': pd.np.float32, 'Species': pd.np.int32})
    test_x, test_y = test, test.pop(y_name)

    return (train_x, train_y), (test_x, test_y)

Define a Keras model (2 hidden layers with 10 neurons in each)


In [7]:
def get_model():
  return keras.Sequential([
    keras.layers.Dense(10, input_shape=(4,), activation=tf.nn.relu, name = "Dense_1"),
    keras.layers.Dense(10, activation=tf.nn.relu, name = "Dense_2"),
    keras.layers.Dense(3, activation=None, name = "logits"),
    keras.layers.Dense(3, activation=tf.nn.softmax, name = "softmax")
  ])

Compiling the model with a distribution strategy

To make the model usable by a TPU, we first must create and compile it using a distribution strategy.


In [8]:
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(TF_MASTER)
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.experimental.TPUStrategy(resolver)


INFO:tensorflow:Initializing the TPU system: grpc://10.110.91.34:8470
INFO:tensorflow:Initializing the TPU system: grpc://10.110.91.34:8470
INFO:tensorflow:Clearing out eager caches
INFO:tensorflow:Clearing out eager caches
INFO:tensorflow:Finished initializing TPU system.
INFO:tensorflow:Finished initializing TPU system.
INFO:tensorflow:Found TPU system:
INFO:tensorflow:Found TPU system:
INFO:tensorflow:*** Num TPU Cores: 8
INFO:tensorflow:*** Num TPU Cores: 8
INFO:tensorflow:*** Num TPU Workers: 1
INFO:tensorflow:*** Num TPU Workers: 1
INFO:tensorflow:*** Num TPU Cores Per Worker: 8
INFO:tensorflow:*** Num TPU Cores Per Worker: 8
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)

In [9]:
with strategy.scope():
  model = get_model()
  model.compile(optimizer=tf.keras.optimizers.Adagrad(learning_rate=0.1), 
                loss=tf.keras.losses.sparse_categorical_crossentropy, 
                metrics=['accuracy'])

model.summary()


Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
Dense_1 (Dense)              (None, 10)                50        
_________________________________________________________________
Dense_2 (Dense)              (None, 10)                110       
_________________________________________________________________
logits (Dense)               (None, 3)                 33        
_________________________________________________________________
softmax (Dense)              (None, 3)                 12        
=================================================================
Total params: 205
Trainable params: 205
Non-trainable params: 0
_________________________________________________________________

Train the model on TPU


In [10]:
# Fetch the data
(train_x, train_y), (test_x, test_y) = load_data()

# Train the model
model.fit(
  train_x.values, train_y.values,
  steps_per_epoch = steps_per_epoch,
  epochs=epochs,
)


Downloading data from http://download.tensorflow.org/data/iris_training.csv
8192/2194 [================================================================================================================] - 0s 0us/step
Downloading data from http://download.tensorflow.org/data/iris_test.csv
8192/573 [============================================================================================================================================================================================================================================================================================================================================================================================================================================] - 0s 0us/step
/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:27: FutureWarning: The pandas.np module is deprecated and will be removed from pandas in a future version. Import numpy directly instead
/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:28: FutureWarning: The pandas.np module is deprecated and will be removed from pandas in a future version. Import numpy directly instead
/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:31: FutureWarning: The pandas.np module is deprecated and will be removed from pandas in a future version. Import numpy directly instead
/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:32: FutureWarning: The pandas.np module is deprecated and will be removed from pandas in a future version. Import numpy directly instead
Epoch 1/50
5/5 [==============================] - 0s 7ms/step - accuracy: 0.4667 - loss: 1.1191
Epoch 2/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.7250 - loss: 0.7990
Epoch 3/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.7000 - loss: 0.5524
Epoch 4/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.7083 - loss: 0.4898
Epoch 5/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.7083 - loss: 0.4448
Epoch 6/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.7417 - loss: 0.4312
Epoch 7/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.7333 - loss: 0.4175
Epoch 8/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9083 - loss: 0.4045
Epoch 9/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.7833 - loss: 0.3834
Epoch 10/50
5/5 [==============================] - 0s 7ms/step - accuracy: 0.8750 - loss: 0.3639
Epoch 11/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.8250 - loss: 0.3659
Epoch 12/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.8750 - loss: 0.3493
Epoch 13/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.8917 - loss: 0.3038
Epoch 14/50
5/5 [==============================] - 0s 7ms/step - accuracy: 0.9250 - loss: 0.2926
Epoch 15/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9333 - loss: 0.2582
Epoch 16/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9250 - loss: 0.2462
Epoch 17/50
5/5 [==============================] - 0s 7ms/step - accuracy: 0.9417 - loss: 0.2132
Epoch 18/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9417 - loss: 0.2183
Epoch 19/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9000 - loss: 0.2533
Epoch 20/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9417 - loss: 0.1799
Epoch 21/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9000 - loss: 0.2484
Epoch 22/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9167 - loss: 0.2163
Epoch 23/50
5/5 [==============================] - 0s 7ms/step - accuracy: 0.9583 - loss: 0.1751
Epoch 24/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9250 - loss: 0.1933
Epoch 25/50
5/5 [==============================] - 0s 7ms/step - accuracy: 0.9667 - loss: 0.1479
Epoch 26/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9500 - loss: 0.1362
Epoch 27/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9833 - loss: 0.1257
Epoch 28/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9667 - loss: 0.1170
Epoch 29/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9667 - loss: 0.1428
Epoch 30/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9750 - loss: 0.1060
Epoch 31/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9750 - loss: 0.1274
Epoch 32/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9667 - loss: 0.1236
Epoch 33/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9583 - loss: 0.0952
Epoch 34/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9750 - loss: 0.0998
Epoch 35/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9750 - loss: 0.0883
Epoch 36/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9667 - loss: 0.1000
Epoch 37/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9750 - loss: 0.0942
Epoch 38/50
5/5 [==============================] - 0s 7ms/step - accuracy: 0.9500 - loss: 0.1247
Epoch 39/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9833 - loss: 0.0843
Epoch 40/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9750 - loss: 0.0798
Epoch 41/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9667 - loss: 0.0967
Epoch 42/50
5/5 [==============================] - 0s 7ms/step - accuracy: 0.9667 - loss: 0.1272
Epoch 43/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9500 - loss: 0.1022
Epoch 44/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9667 - loss: 0.0886
Epoch 45/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9833 - loss: 0.0767
Epoch 46/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9833 - loss: 0.0804
Epoch 47/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9833 - loss: 0.0850
Epoch 48/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9750 - loss: 0.0906
Epoch 49/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9833 - loss: 0.0989
Epoch 50/50
5/5 [==============================] - 0s 6ms/step - accuracy: 0.9667 - loss: 0.0911
Out[10]:
<tensorflow.python.keras.callbacks.History at 0x7f1c7383cc88>

Evaluation of the model


In [11]:
model.evaluate(test_x.values, test_y.values,
    batch_size=8)


4/4 [==============================] - 0s 10ms/step - accuracy: 0.9333 - loss: 0.0730
Out[11]:
[0.07300268858671188, 0.9333333969116211]

Save the model


In [12]:
model.save_weights('./DNN_TPU_1024.h5', overwrite=True)

Prediction

Prediction data


In [13]:
COLUMNS_NAME=['SepalLength', 'SepalWidth', 'PetalLength', 'PetalWidth']
data = pd.DataFrame(PREDICTION_INPUT_DATA, columns=COLUMNS_NAME)
print(data)


   SepalLength  SepalWidth  PetalLength  PetalWidth
0          6.9         3.1          5.4         2.1
1          5.1         3.3          1.7         0.5
2          5.9         3.0          4.2         1.5
3          6.0         3.4          4.5         1.6
4          5.5         2.5          4.0         1.3
5          6.2         2.9          4.3         1.3
6          5.5         4.2          1.4         0.2
7          6.3         2.8          5.1         1.5

Prediction on TPU


In [14]:
predictions = model.predict(data.values.astype(np.float32))
template = ('\nPrediction is "{}" ({:.1f}%), expected "{}"')
for pred_dict, expec in zip(predictions, PREDICTION_OUTPUT_DATA):
  class_index = np.argmax(pred_dict)
  class_probability = np.max(pred_dict)
  print(template.format(SPECIES[class_index], 100*class_probability, expec))


Prediction is "Virginica" (87.8%), expected "Virginica"

Prediction is "Setosa" (99.8%), expected "Setosa"

Prediction is "Versicolor" (97.3%), expected "Versicolor"

Prediction is "Versicolor" (97.3%), expected "Versicolor"

Prediction is "Versicolor" (97.2%), expected "Versicolor"

Prediction is "Versicolor" (97.3%), expected "Versicolor"

Prediction is "Setosa" (100.0%), expected "Setosa"

Prediction is "Versicolor" (55.7%), expected "Virginica"

Prediction on CPU


In [15]:
cpu_model = get_model()
cpu_model.load_weights('./DNN_TPU_1024.h5')
cpu_predictions = cpu_model.predict(data)
template = ('\nPrediction is "{}" ({:.1f}%), expected "{}"')
for pred_dict, expec in zip(cpu_predictions, PREDICTION_OUTPUT_DATA):
  class_index = np.argmax(pred_dict)
  class_probability = np.max(pred_dict)
  print(template.format(SPECIES[class_index], 100*class_probability, expec))


Prediction is "Virginica" (87.2%), expected "Virginica"

Prediction is "Setosa" (99.8%), expected "Setosa"

Prediction is "Versicolor" (97.2%), expected "Versicolor"

Prediction is "Versicolor" (97.2%), expected "Versicolor"

Prediction is "Versicolor" (97.2%), expected "Versicolor"

Prediction is "Versicolor" (97.2%), expected "Versicolor"

Prediction is "Setosa" (100.0%), expected "Setosa"

Prediction is "Versicolor" (54.7%), expected "Virginica"

What's next

  • Learn about Cloud TPUs that Google designed and optimized specifically to speed up and scale up ML workloads for training and inference and to enable ML engineers and researchers to iterate more quickly.
  • Explore the range of Cloud TPU tutorials and Colabs to find other examples that can be used when implementing your ML project.

On Google Cloud Platform, in addition to GPUs and TPUs available on pre-configured deep learning VMs, you will find AutoML(beta) for training custom models without writing code and Cloud ML Engine which will allows you to run parallel trainings and hyperparameter tuning of your custom models on powerful distributed hardware.