つれづれなる備忘録

日々の発見をあるがままに綴る

トップ > その他 > Google Colabでpython3 ～ TPUの利用

Google Colabでpython3 ～ TPUの利用

その他 ツール python Google Colab

　前回はGoogle ColabでGPUの使用方法について紹介したが、今回はもう一つのアクセラレータであるTPUの使用方法と効果について紹介する。

atatat.hatenablog.com

1. TPU

　TPUとはTensor Processing Unitの略でGoogleがディープラーニングなど機械学習用に開発した集積回路(ASIC)で既存のCPU/GPUに対して演算速度向上とともに消費電力を抑えることが特徴となる。 Cloud TPU v3 Podの演算能力は100PFLOPS以上となっており、適用できる用途は限られるものの近年のスパコン並みの性能を有している。(詳しくはCloud TPU  |  Google Cloud)

2. TPUの設定

　Google Colabのトップページの下のアクセラレターの使用：TensorFlow と TPUのページ(英語）に従って設定をする。以前紹介したGPUの設定と同様にノートブックのメニューバー上の"編集"から"ノートブックの設定"をクリックすると、以下のアクセラレータを選択するポップアップがあらわれるのでTPUを選択して保存する。

"TPUの設定"
TPUの設定

TPUが設定できたかどうか確認するには以下のコードを実行する。

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) # new command

なお最終行(New commandのところ)はもとのページにはtpu_strategy = tf.distribute.experimental.TPUStrategy(tpu)となっていたが、実行するとwarningが出てtf.distribute.TPUStrategyを使うようにメッセージがあらわれるのでtpu_strategy = tf.distribute.TPUStrategy(tpu)とした。正しく設定できていれば以下のように出力される。

Running on TPU  [' IP address']
INFO:tensorflow:Initializing the TPU system: grpc://IP address
INFO:tensorflow:Initializing the TPU system: grpc://IP address
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)

3. TPUの機械学習例

　"TensorFlow と TPU"のページには花の画像と名前を画像識別する例があるがコードが少し長いので、一番下にあるFashion MNISTの例を紹介する。Fashion MNISTとは手書き数字認識MNISTと同じフォーマット28x28の10種類の手書き数字画像を、10種類の衣料品の画像識別に置き換えたものである。識別されるものは、0 T-シャツ/トップ (T-shirt/top), 1 ズボン (Trouser), 2 プルオーバー (Pullover), 3 ドレス (Dress), 4 コート (Coat), 5 サンダル (Sandal), 6 シャツ (Shirt), 7 スニーカー (Sneaker), 8 バッグ (Bag), 9 アンクルブーツ (Ankle boot)

Fashion MNISTのラベル付き画像データセットは以下のコードでロードし、画像はx_train,x_test、ラベルはy_train,y_testにデータが格納される。

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)

実行すると以下のように表示される。

Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/train-labels-idx1-ubyte.gz
32768/29515 [=================================] - 0s 0us/step
Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/train-images-idx3-ubyte.gz
26427392/26421880 [==============================] - 0s 0us/step
Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/t10k-labels-idx1-ubyte.gz
8192/5148 [===============================================] - 0s 0us/step
Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/t10k-images-idx3-ubyte.gz
4423680/4422102 [==============================] - 0s 0us/step

次に学習に用いるディープラーニングのモデルを下記のように定義する。

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

TPUを用いた学習を実行するために以下のコードを実行する。

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
)

model.save_weights('./fashion_mnist.h5', overwrite=True)

実行すると以下のように学習の過程が表示される。バッチごとに精度が上がっていき最初の精度が0.66から最後は0.94になっている。最初と最後は少し時間がかかっているが、途中は1バッチ当たりの実行速度は平均1秒程度で恐らく驚異的に早い。

Epoch 1/17
 1/60 [..............................] - ETA: 2:43 - loss: 3.8517 - sparse_categorical_accuracy: 0.1000WARNING:tensorflow:Callbacks method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0022s vs `on_train_batch_end` time: 0.0200s). Check your callbacks.
WARNING:tensorflow:Callbacks method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0022s vs `on_train_batch_end` time: 0.0200s). Check your callbacks.
60/60 [==============================] - 4s 68ms/step - loss: 1.2234 - sparse_categorical_accuracy: 0.6661
Epoch 2/17
60/60 [==============================] - 1s 21ms/step - loss: 0.5432 - sparse_categorical_accuracy: 0.8137
Epoch 3/17
60/60 [==============================] - 1s 21ms/step - loss: 0.4377 - sparse_categorical_accuracy: 0.8472
Epoch 4/17
60/60 [==============================] - 1s 22ms/step - loss: 0.3777 - sparse_categorical_accuracy: 0.8672
Epoch 5/17
60/60 [==============================] - 1s 22ms/step - loss: 0.3462 - sparse_categorical_accuracy: 0.8794
Epoch 6/17
60/60 [==============================] - 1s 22ms/step - loss: 0.3200 - sparse_categorical_accuracy: 0.8860
Epoch 7/17
60/60 [==============================] - 1s 21ms/step - loss: 0.2893 - sparse_categorical_accuracy: 0.8961
Epoch 8/17
60/60 [==============================] - 1s 21ms/step - loss: 0.2751 - sparse_categorical_accuracy: 0.9007
Epoch 9/17
60/60 [==============================] - 1s 22ms/step - loss: 0.2542 - sparse_categorical_accuracy: 0.9072
Epoch 10/17
60/60 [==============================] - 1s 22ms/step - loss: 0.2320 - sparse_categorical_accuracy: 0.9158
Epoch 11/17
60/60 [==============================] - 1s 22ms/step - loss: 0.2224 - sparse_categorical_accuracy: 0.9175
Epoch 12/17
60/60 [==============================] - 1s 22ms/step - loss: 0.2074 - sparse_categorical_accuracy: 0.9228
Epoch 13/17
60/60 [==============================] - 1s 21ms/step - loss: 0.1988 - sparse_categorical_accuracy: 0.9272
Epoch 14/17
60/60 [==============================] - 1s 21ms/step - loss: 0.1849 - sparse_categorical_accuracy: 0.9312
Epoch 15/17
60/60 [==============================] - 1s 22ms/step - loss: 0.1806 - sparse_categorical_accuracy: 0.9319
Epoch 16/17
60/60 [==============================] - 1s 22ms/step - loss: 0.1706 - sparse_categorical_accuracy: 0.9364
Epoch 17/17
58/60 [============================>.] - ETA: 0s - loss: 0.1593 - sparse_categorical_accuracy: 0.9409WARNING: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.0132s). 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.0132s). Check your callbacks.
60/60 [==============================] - 7s 124ms/step - loss: 0.1600 - sparse_categorical_accuracy: 0.9407 - val_loss: 0.2365 - val_sparse_categorical_accuracy: 0.9238

4. 画像識別例

　学習結果に基づいて画像の識別例について示す。識別を実行する関数を以下のように定義する。

LABEL_NAMES = ['t_shirt', 'trouser', 'pullover', 'dress', 'coat', 'sandal', 'shirt', 'sneaker', 'bag', 'ankle_boots']

cpu_model = create_model()
cpu_model.load_weights('./fashion_mnist.h5')

from matplotlib import pyplot
%matplotlib inline

def plot_predictions(images, predictions):
  n = images.shape[0]
  nc = int(np.ceil(n / 4))
  f, axes = pyplot.subplots(nc, 4)
  for i in range(nc * 4):
    y = i // 4
    x = i % 4
    axes[x, y].axis('off')
    
    label = LABEL_NAMES[np.argmax(predictions[i])]
    confidence = np.max(predictions[i])
    if i > n:
      continue
    axes[x, y].imshow(images[i])
    axes[x, y].text(0.5, 0.5, label + '\n%.3f' % confidence, fontsize=14)

15番までのテスト画像に対して識別を実行してみる。（1部うまく実行できていない部分があるが、そのまま表示する）

pyplot.gcf().set_size_inches(8, 8)  
plot_predictions(np.squeeze(x_test[:15]), cpu_model.predict(x_test[:15]))

少し見にくいがテスト画像の上に識別（予測）結果と確定度を表示している。

"Fashion MNIST識別テスト"
Fashion MNIST識別テスト

5. まとめ

　今回は機械学習・ディープラーニングに特化したアクセラレータとしてTPUの利用方法と使用例について紹介した。TPUを上手に利用することで、画像識別などの時間がかかる学習時間を格段に短縮することができる。