Google ColabのTPUをGPUと比較してみた

単純なコードでTPUとGPUを比較してみたいと思ったので。
[Google Colabでpython3 ～ TPUの利用](https://atatat.hatenablog.com/entry/colab_python3_tpu)を参考にしました。

#GPUのコード
```python
import tensorflow as tf
import numpy as np
import os
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.fashion_mnist.load_data()

# add empty color dimension
x_train = np.expand_dims(x_train, -1)
x_test = np.expand_dims(x_test, -1)

def create_model():
  model = tf.keras.models.Sequential()
  model.add(tf.keras.layers.BatchNormalization(input_shape=x_train.shape[1:]))
  model.add(tf.keras.layers.Conv2D(64, (5, 5), padding='same', activation='elu'))
  model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=(2,2)))
  model.add(tf.keras.layers.Dropout(0.25))

  model.add(tf.keras.layers.BatchNormalization(input_shape=x_train.shape[1:]))
  model.add(tf.keras.layers.Conv2D(128, (5, 5), padding='same', activation='elu'))
  model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2)))
  model.add(tf.keras.layers.Dropout(0.25))

  model.add(tf.keras.layers.BatchNormalization(input_shape=x_train.shape[1:]))
  model.add(tf.keras.layers.Conv2D(256, (5, 5), padding='same', activation='elu'))
  model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=(2,2)))
  model.add(tf.keras.layers.Dropout(0.25))

  model.add(tf.keras.layers.Flatten())
  model.add(tf.keras.layers.Dense(256))
  model.add(tf.keras.layers.Activation('elu'))
  model.add(tf.keras.layers.Dropout(0.5))
  model.add(tf.keras.layers.Dense(10))
  model.add(tf.keras.layers.Activation('softmax'))
  return model

model = create_model()
model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3, ),
    loss='sparse_categorical_crossentropy',
    metrics=['sparse_categorical_accuracy'])

model.fit(
    x_train.astype(np.float32), y_train.astype(np.float32),
    epochs=17,
    steps_per_epoch=60,
    validation_data=(x_test.astype(np.float32), y_test.astype(np.float32)),
    validation_freq=17
)
```
これを実行するとこんな感じ。

```
Epoch 1/17
 2/60 [>.............................] - ETA: 1s - loss: 5.3304 - sparse_categorical_accuracy: 0.1695WARNING:tensorflow:Callbacks method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0066s vs `on_train_batch_end` time: 0.0533s). Check your callbacks.
60/60 [==============================] - 4s 59ms/step - loss: 1.1856 - sparse_categorical_accuracy: 0.6706
Epoch 2/17
60/60 [==============================] - 4s 60ms/step - loss: 0.5417 - sparse_categorical_accuracy: 0.8126
Epoch 3/17
60/60 [==============================] - 4s 60ms/step - loss: 0.4465 - sparse_categorical_accuracy: 0.8463
Epoch 4/17
60/60 [==============================] - 4s 61ms/step - loss: 0.3904 - sparse_categorical_accuracy: 0.8628
Epoch 5/17
60/60 [==============================] - 4s 60ms/step - loss: 0.3479 - sparse_categorical_accuracy: 0.8758
Epoch 6/17
60/60 [==============================] - 4s 61ms/step - loss: 0.3194 - sparse_categorical_accuracy: 0.8853
Epoch 7/17
60/60 [==============================] - 4s 60ms/step - loss: 0.2883 - sparse_categorical_accuracy: 0.8961
Epoch 8/17
60/60 [==============================] - 4s 60ms/step - loss: 0.2715 - sparse_categorical_accuracy: 0.9024
Epoch 9/17
60/60 [==============================] - 4s 60ms/step - loss: 0.2492 - sparse_categorical_accuracy: 0.9093
Epoch 10/17
60/60 [==============================] - 4s 60ms/step - loss: 0.2381 - sparse_categorical_accuracy: 0.9137
Epoch 11/17
60/60 [==============================] - 4s 60ms/step - loss: 0.2245 - sparse_categorical_accuracy: 0.9168
Epoch 12/17
60/60 [==============================] - 4s 60ms/step - loss: 0.2092 - sparse_categorical_accuracy: 0.9231
Epoch 13/17
60/60 [==============================] - 4s 61ms/step - loss: 0.2013 - sparse_categorical_accuracy: 0.9248
Epoch 14/17
60/60 [==============================] - 4s 60ms/step - loss: 0.1909 - sparse_categorical_accuracy: 0.9298
Epoch 15/17
60/60 [==============================] - 4s 61ms/step - loss: 0.1805 - sparse_categorical_accuracy: 0.9325
Epoch 16/17
60/60 [==============================] - 4s 60ms/step - loss: 0.1698 - sparse_categorical_accuracy: 0.9371
Epoch 17/17
60/60 [==============================] - 5s 81ms/step - loss: 0.1609 - sparse_categorical_accuracy: 0.9396 - val_loss: 0.2447 - val_sparse_categorical_accuracy: 0.9141
```

#TPUのコード
ほとんど同じなのだが、あえて全部載せる。

```python
import tensorflow as tf
try:
  tpu = tf.distribute.cluster_resolver.TPUClusterResolver()  # TPU detection
  print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])
except ValueError:
  raise BaseException('ERROR: Not connected to a TPU runtime; please see the previous cell in this notebook for instructions!')

tf.config.experimental_connect_to_cluster(tpu)
tf.tpu.experimental.initialize_tpu_system(tpu)
tpu_strategy = tf.distribute.TPUStrategy(tpu)

import numpy as np
import os
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.fashion_mnist.load_data()

# add empty color dimension
x_train = np.expand_dims(x_train, -1)
x_test = np.expand_dims(x_test, -1)

def create_model():
  model = tf.keras.models.Sequential()
  model.add(tf.keras.layers.BatchNormalization(input_shape=x_train.shape[1:]))
  model.add(tf.keras.layers.Conv2D(64, (5, 5), padding='same', activation='elu'))
  model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=(2,2)))
  model.add(tf.keras.layers.Dropout(0.25))

  model.add(tf.keras.layers.BatchNormalization(input_shape=x_train.shape[1:]))
  model.add(tf.keras.layers.Conv2D(128, (5, 5), padding='same', activation='elu'))
  model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2)))
  model.add(tf.keras.layers.Dropout(0.25))

  model.add(tf.keras.layers.BatchNormalization(input_shape=x_train.shape[1:]))
  model.add(tf.keras.layers.Conv2D(256, (5, 5), padding='same', activation='elu'))
  model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=(2,2)))
  model.add(tf.keras.layers.Dropout(0.25))

  model.add(tf.keras.layers.Flatten())
  model.add(tf.keras.layers.Dense(256))
  model.add(tf.keras.layers.Activation('elu'))
  model.add(tf.keras.layers.Dropout(0.5))
  model.add(tf.keras.layers.Dense(10))
  model.add(tf.keras.layers.Activation('softmax'))
  return model

tf.keras.backend.clear_session()

resolver = tf.distribute.cluster_resolver.TPUClusterResolver('grpc://' + os.environ['COLAB_TPU_ADDR'])
tf.config.experimental_connect_to_cluster(resolver)

# This is the TPU initialization code that has to be at the beginning.
tf.tpu.experimental.initialize_tpu_system(resolver)
print("All devices: ", tf.config.list_logical_devices('TPU'))

strategy = tf.distribute.experimental.TPUStrategy(resolver)

with strategy.scope():
  model = create_model()
  model.compile(
      optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3, ),
      loss='sparse_categorical_crossentropy',
      metrics=['sparse_categorical_accuracy'])

model.fit(
    x_train.astype(np.float32), y_train.astype(np.float32),
    epochs=17,
    steps_per_epoch=60,
    validation_data=(x_test.astype(np.float32), y_test.astype(np.float32)),
    validation_freq=17
)
```
これを実行するとこんな感じに。

```
Epoch 1/17
 1/60 [..............................] - ETA: 2:47 - loss: 3.6474 - sparse_categorical_accuracy: 0.0970WARNING:tensorflow:Callbacks method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0023s vs `on_train_batch_end` time: 0.0206s). Check your callbacks.
WARNING:tensorflow:Callbacks method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0023s vs `on_train_batch_end` time: 0.0206s). Check your callbacks.
60/60 [==============================] - 4s 68ms/step - loss: 1.1290 - sparse_categorical_accuracy: 0.6805
Epoch 2/17
60/60 [==============================] - 1s 22ms/step - loss: 0.5311 - sparse_categorical_accuracy: 0.8190
Epoch 3/17
60/60 [==============================] - 1s 21ms/step - loss: 0.4299 - sparse_categorical_accuracy: 0.8518
Epoch 4/17
60/60 [==============================] - 1s 22ms/step - loss: 0.3816 - sparse_categorical_accuracy: 0.8665
Epoch 5/17
60/60 [==============================] - 1s 22ms/step - loss: 0.3355 - sparse_categorical_accuracy: 0.8811
Epoch 6/17
60/60 [==============================] - 1s 21ms/step - loss: 0.3122 - sparse_categorical_accuracy: 0.8883
Epoch 7/17
60/60 [==============================] - 1s 22ms/step - loss: 0.2810 - sparse_categorical_accuracy: 0.8977
Epoch 8/17
60/60 [==============================] - 1s 22ms/step - loss: 0.2651 - sparse_categorical_accuracy: 0.9033
Epoch 9/17
60/60 [==============================] - 1s 22ms/step - loss: 0.2389 - sparse_categorical_accuracy: 0.9133
Epoch 10/17
60/60 [==============================] - 1s 22ms/step - loss: 0.2345 - sparse_categorical_accuracy: 0.9151
Epoch 11/17
60/60 [==============================] - 1s 22ms/step - loss: 0.2146 - sparse_categorical_accuracy: 0.9213
Epoch 12/17
60/60 [==============================] - 1s 22ms/step - loss: 0.1992 - sparse_categorical_accuracy: 0.9261
Epoch 13/17
60/60 [==============================] - 1s 22ms/step - loss: 0.1895 - sparse_categorical_accuracy: 0.9292
Epoch 14/17
60/60 [==============================] - 1s 22ms/step - loss: 0.1813 - sparse_categorical_accuracy: 0.9331
Epoch 15/17
60/60 [==============================] - 1s 22ms/step - loss: 0.1682 - sparse_categorical_accuracy: 0.9366
Epoch 16/17
60/60 [==============================] - 1s 22ms/step - loss: 0.1634 - sparse_categorical_accuracy: 0.9385
Epoch 17/17
60/60 [==============================] - ETA: 0s - loss: 0.1548 - sparse_categorical_accuracy: 0.9415WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/data/ops/multi_device_iterator_ops.py:601: get_next_as_optional (from tensorflow.python.data.ops.iterator_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Use `tf.data.Iterator.get_next_as_optional()` instead.
WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/data/ops/multi_device_iterator_ops.py:601: get_next_as_optional (from tensorflow.python.data.ops.iterator_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Use `tf.data.Iterator.get_next_as_optional()` instead.
WARNING:tensorflow:Callbacks method `on_test_batch_end` is slow compared to the batch time (batch time: 0.0016s vs `on_test_batch_end` time: 0.0114s). Check your callbacks.
WARNING:tensorflow:Callbacks method `on_test_batch_end` is slow compared to the batch time (batch time: 0.0016s vs `on_test_batch_end` time: 0.0114s). Check your callbacks.
60/60 [==============================] - 7s 121ms/step - loss: 0.1548 - sparse_categorical_accuracy: 0.9415 - val_loss: 0.2249 - val_sparse_categorical_accuracy: 0.9239
```
は？？？4倍速くなってる！TPUすごい！
# 参考記事
* [Google Colabでpython3 ～ TPUの利用](https://atatat.hatenablog.com/entry/colab_python3_tpu)
* [Google ColaboratoryのTPUを試してみる](https://qiita.com/koshian2/items/25a6341c035e8a260a01)
* [TensorFlow 2.0でTPUを利用する](https://techblog.gmo-ap.jp/2020/06/09/tensorflow-2-0-tpu/)
* [TensorFlow2.0 with KerasでいろいろなMNIST（TPU対応）](https://qiita.com/koshian2/items/5029028fe6d01b243f44)