Google ColaboratoryでTPUとKerasを使う

#これはなに？
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](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/240237/4dbdf777-2341-3edb-b4ab-3995fc8c0ba6.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