pdfcoffee

Chapter 16
Using TensorFlow 2.1 and nightly build
As of November 2019, you can get full TPU support only with the latest TensorFlow
2.x nightly build. If you use the Google Cloud Console (https://console.cloud.
google.com/) you can get the latest nightly build. Just, go to Compute Engine
| TPUs | CREATE TPU NODE. The version selector has a "nightly-2.x" option.
Martin Görner has a nice demo at http://bit.ly/keras-tpu-tf21 (see Figure 13).
This is used for classifying images of flowers:
Figure 13: Martin Görner on Twitter on Full Keras/TPU support
Note that both Regular Keras using model.fit() and custom training loop,
distributed are supported. You can refer to http://bit.ly/keras-tpu-tf21. Let's
look at the most important parts of the code related to TPUs. First at all, the imports:
import re
import tensorflow as tf
import numpy as np
from matplotlib import pyplot as plt
print("Tensorflow version " + tf.__version__)
Then the detection of the TPUs, and the selection of TPU strategy:
try:
tpu = tf.distribute.cluster_resolver.TPUClusterResolver()
# TPU detection
print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])
except ValueError:
tpu = None
if tpu:
tf.config.experimental_connect_to_cluster(tpu)
tf.tpu.experimental.initialize_tpu_system(tpu)
strategy = tf.distribute.experimental.TPUStrategy(tpu)
else:
strategy = tf.distribute.get_strategy()
print("REPLICAS: ", strategy.num_replicas_in_sync)
[ 587 ]