Go to Qiita Advent Calendar 2024 Top
LoginSignup
13
6

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

More than 3 years have passed since last update.

@Fulltea

Google ColaboratoryでTPUとKerasを使う

Last updated at Posted at 2019-12-07

#これはなに？
Google ColaboratoryでTPUを使う記事です。
ランタイムを切り替えるだけで動作するGPUと異なり、いくつかコードに書き足す点があったため、備忘録として記しておきます。

#環境
Google Colaboratoryです
tensorflow 1.15.0がインストールされています。

TPUtest.py
import tensorflow as tf
import distuitls

print(distutils.version.LooseVersion(tf.__version__))
#>>1.15.0

#検証コード
CNNを用いて、mnistを分類します。
Googleによると、TPUはCNNに最適化されていないらしいので、若干TPUに不利な条件かもしれません。が、性能評価を厳密にやりたいわけではないのでまあいいでしょう。

##データの準備、加工

TPUtest.py
from tensorflow.keras.datasets import mnist
from tensorflow.keras.utils import to_categorical
import numpy as np

#データダウンロード
(X_train, y_train), (X_test, y_test) = mnist.load_data()

#255で割る
X_train = X_train/255
X_test = X_test/255

#画像データの形状を変更
X_train = X_train.reshape(-1,28,28,1).astype(np.float32)
X_test = X_test.reshape(-1,28,28,1).astype(np.float32)

#正解ラベルをone-hotに変換
y_train = to_categorical(y_train)
y_test = to_categorical(y_test)

##モデルの構築、コンパイル

TPUtest.py
from tensorflow.keras.layers import Conv2D, Dense, ReLU, Flatten, Input, MaxPool2D, Dropout
from tensorflow.keras.models import Sequential
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.utils import plot_model

def getModel():
  model = Sequential()
  model.add(Conv2D(3,3,input_shape=(28,28,1)))
  model.add(MaxPool2D(2))
  model.add(ReLU())
  model.add(Dropout(0.2))
  model.add(Flatten())
  model.add(Dense(1024))
  model.add(ReLU())
  model.add(Dense(10, activation="softmax"))
  return model

model = getModel()

#描画
plot_model(model, show_shapes=True, show_layer_names=False)

#コンパイル
model.compile(Adam(), loss="categorical_crossentropy", metrics=["acc"])

model.png

至って普通のモデルです。

##訓練

TPUtest.py
%%time
model.fit(X_train, y_train, epochs=10, validation_data=(X_test,y_test))

##予測

TPUtest.py
%%time
y_pred = model.predict(X_test)

##予測結果の表示

TPUtest.py
from sklearn.metrics import accuracy_score
import numpy as np

#one-hotベクトルをもとに戻す
y_pred = np.argmax(y_pred, axis=1)
y_test = np.argmax(y_test, axis=1)
print(accuracy_score(y_pred, y_test))
#>>0.9854

このコードを3種類のランタイムで実行して比較します。

#実行結果
実行結果は以下のとおりです。

ランタイム 訓練時間 予測時間 予測スコア
CPU 37s/epoch 1.49s 0.9854
GPU 13s/epoch 0.54s 0.9859
TPU 37s/epoch 2.76s 0.9863

...TPU動いてなくない?

#TPUを働かせる
まずはデバイスを確認します

TPUtest.py
import os
import tensorflow as tf
import pprint

if 'COLAB_TPU_ADDR' not in os.environ:
  print('ERROR: Not connected to a TPU runtime; please see the first cell in this notebook for instructions!')
else:
  tpu_address = 'grpc://' + os.environ['COLAB_TPU_ADDR']
  print ('TPU address is', tpu_address)

  with tf.Session(tpu_address) as session:
    devices = session.list_devices()
    
  print('TPU devices:')
  pprint.pprint(devices)

ずらずらっと表示されればOKです。

##TPU向けのコンパイル
モデルの作成、コンパイルの際に少しだけ工夫が必要です。

TPUtest.py
def getModel():
  # ここまではGPUと変わらないため省略
  return model

# TPUの各種セットアップ
resolver = tf.contrib.cluster_resolver.TPUClusterResolver('grpc://' + os.environ['COLAB_TPU_ADDR'])
tf.contrib.distribute.initialize_tpu_system(resolver)
strategy = tf.contrib.distribute.TPUStrategy(resolver)


with strategy.scope():# これを追記する必要がある
  model = getModel()
  model.compile(Adam(), loss="categorical_crossentropy", metrics=["acc"])

#あとは普通にfitする
model.fit(X_train, y_train, epochs=10, validation_data=(X_test,y_test))

CPUのときより明らかに快適に学習が進みます。
そのまま予測も試しましょう。

TPUtest.py
y_pred = model.predict(X_test)

え？動くけどめっちゃ遅くない...?
予測終わらなくない??
実行結果は以下の通りです。

ランタイム 訓練時間 予測時間 予測スコア
CPU(再掲) 37s/epoch 1.49s 0.9854
GPU(再掲) 13s/epoch 0.54s 0.9859
TPU 17.7s/epoch 15min 15s 0.9853

##予測をまともな速度で行う
TPUで予測を行うととんでもない速度になることがわかりました。validationは(まだ)常識的な速度で行っているのに、なんでだ...?

やむをえないので、学習はTPUで行ったあと、予測はCPUで行います。

TPUtest.py
# TPUでの学習
model.fit(X_train, y_train, epochs=10, validation_data=(X_test,y_test))
model.save_weights("./weight.h5")# 重みをファイルに保存

#CPUでの予測
cpu_model = getModel()# CPUでモデルを構築
cpu_model.load_weights("./weight.h5")# 保存した重みを読み込む
y_pred = cpu_model.predict(X_test)# cpu_modelで予測

最終的な性能は以下の通り。

ランタイム 訓練時間 予測時間 予測スコア
CPU 37s/epoch 1.49s 0.9854
GPU 13s/epoch 0.54s 0.9859
TPU 17.7s/epoch 1.22s(CPUを使用) 0.9853

正直GPUのほうが楽かなって気もしますが、先述したとおりCNNはTPUの苦手分野らしく、例えばLSTMはGPUの倍以上の学習速度が出たので、状況に応じて使い分けてもいいかもしれません。単純にランタイムを2個同時に動かせるし。

#ハマった点
エラーにそこそこ遭遇しました...
##InvalidArgumentError その1
エラーメッセージ

InvalidArgumentError: Cannot assign a device for operation conv2d_1/kernel/IsInitialized/VarIsInitializedOp: node conv2d_1/kernel/IsInitialized/VarIsInitializedOp (defined at /usr/local/lib/python3.6/dist-packages/tensorflow_core/python/framework/ops.py:1748)  was explicitly assigned to /job:worker/replica:0/task:0/device:TPU:0 but available devices are [ /job:localhost/replica:0/task:0/device:CPU:0, /job:localhost/replica:0/task:0/device:XLA_CPU:0 ]. Make sure the device specification refers to a valid device.
	 [[conv2d_1/kernel/IsInitialized/VarIsInitializedOp]]

kerasを用いてmodelを作成するとエラーが出ました。
kerasではなくtensorflow.kerasを用いることで動きました。罠でしょ。

##InvalidArgumentError その2
エラーメッセージ

InvalidArgumentError: Unsupported data type for TPU: double, caused by output IteratorGetNext:0

TPUはDouble型に対応していないようなので、np.float32に変換してから学習させます。

##InvalidArgumentError その3
エラーメッセージ

InvalidArgumentError: No OpKernel was registered to support Op 'TPUReplicatedInput' used by node input0_1 (defined at /usr/local/lib/python3.6/dist-packages/tensorflow_core/python/framework/ops.py:1748) with these attrs: [T=DT_INT32, N=8]
Registered devices: [CPU, XLA_CPU]
Registered kernels:
  <no registered kernels>

	 [[input0_1]]

偶然に一度だけ発生したエラー。とりあえず再起動したら動いた。
再現性の検証とかはしたくないのでやってないです。

##InvalidArgumentError その4
エラーメッセージ

InvalidArgumentError: Cannot assign a device for operation lstm_1/random_uniform/RandomUniform: node lstm_1/random_uniform/RandomUniform (defined at /usr/local/lib/python3.6/dist-packages/tensorflow_core/python/framework/ops.py:1748)  was explicitly assigned to /job:worker/replica:0/task:0/device:TPU:0 but available devices are [ /job:localhost/replica:0/task:0/device:CPU:0, /job:localhost/replica:0/task:0/device:XLA_CPU:0 ]. Make sure the device specification refers to a valid device.
	 [[lstm_1/random_uniform/RandomUniform]]

TPUにアクセスできない...?とりあえず再起動。

##InternalError
エラーメッセージ

InternalError: Failed to serialize message

LSTMに大量のデータを読み込ませたら発生しました。量を減らしたら動いてくれました。
メモリエラー、なのでしょうか?
(同じ量のデータをGPUランタイムに渡すと動くので謎です。まあGPUでも処理時間が長くなりすぎて12hで終わらないので意味はなかったんですけど)

##KeyError
エラーメッセージ

KeyError: 'COLAB_TPU_ADDR'

ランタイムがTPUではないときに発生します。
TPUに切り替えて実行してください。

#参考文献
https://colab.research.google.com/github/tensorflow/tpu/blob/master/tools/colab/fashion_mnist.ipynb#scrollTo=2a5cGsSTEBQD

13
6
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
13
6

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?