pdfcoffee

Tensor Processing Unit
Note that full support for TPUDistributionStrategy is planned
for TensorFlow 2.1. 2.0 has limited support, and the issue is
tracked in https://github.com/tensorflow/tensorflow/
issues/24412.
So, let's define a standard CNN model made up of three convolutional layers,
alternated with max-pooling layers and followed by two dense layers with a dropout
in the middle. For the sake of brevity, the definition of input_shape, batch_size is
omitted. In this case, we use the functional tf.keras API (see Chapter 2, TensorFlow
1.x and 2.x):
Inp = tf.keras.Input(name='input', shape=input_shape, batch_
size=batch_size, dtype=tf.float32)
x = Conv2D(32, kernel_size=(3, 3), activation='relu',name = 'Conv_01')
(Inp)
x = MaxPooling2D(pool_size=(2, 2),name = 'MaxPool_01')(x)
x = Conv2D(64, (3, 3), activation='relu',name = 'Conv_02')(x)
x = MaxPooling2D(pool_size=(2, 2),name = 'MaxPool_02')(x)
x = Conv2D(64, (3, 3), activation='relu',name = 'Conv_03')(x)
x = Flatten(name = 'Flatten_01')(x)
x = Dense(64, activation='relu',name = 'Dense_01')(x)
x = Dropout(0.5,name = 'Dropout_02')(x)
output = Dense(num_classes, activation='softmax',name = 'Dense_02')(x)
model = tf.keras.Model(inputs=[Inp], outputs=[output])
Let's now use Adam optimizer and compile the model:
#Use a tf optimizer rather than a Keras one for now
opt = tf.train.AdamOptimizer(learning_rate)
model.compile(
optimizer=opt,
loss='categorical_crossentropy',
metrics=['acc'])
Then, we call tpu.keras_to_tpu_model to convert to a TPU model and then we use
the tpu.TPUDistributionStrategy for running on a TPU. It's as simple as that; we
just need to take the appropriate strategy with TPUDistributionStrategy() and all
the rest is done transparently on our behalf:
tpu_model = tf.contrib.tpu.keras_to_tpu_model(
model,
strategy=tf.contrib.tpu.TPUDistributionStrategy(
tf.contrib.cluster_resolver.TPUClusterResolver(TPU_ADDRESS)))
[ 582 ]