Google ColabのTPUで対GPUの最速に挑戦する

[前回の記事](https://qiita.com/koshian2/items/25a6341c035e8a260a01)が思わぬ反響で驚愕していますが、今回はColabのTPUを限界まで試してみたいと思います。特殊な条件にするとColabのTPUは、GPU比で20倍以上速くなることがわかりました。しかも無料です。それを見ていきましょう。

## TPUがGPUと比べて速くなる条件とは
前回の結果からTPUがGPUと比べて速くなる条件は2つあることがわかりました。

* 多層パーセプトロン（MLP）よりも畳み込みニューラルネットワーク（CNN）であること
* バッチサイズを大きくできること

1点目ですが、CNNだから速いのではなく、CNNがパラメーター数の割に計算量が多い、つまり**訓練時間のうち計算量が支配的になるから**と捉えることもできます。なぜなら、CNNは1つの畳み込みフィルターでパラメーターを使いまわしするのですから。

また[GPUで行ったフレームワーク別の調査をTPUに適用したところ](https://github.com/koshian2/ResNet-MultipleFramework)、TPUではモデル（ResNet）が小さすぎて、層を深くしてもほぼ計算時間が変わらないという現象がありました[^1]。

![](https://raw.githubusercontent.com/koshian2/ResNet-MultipleFramework/master/images/times_tpu.png)

[^1]: GPUのバッチサイズは128で統一したのに対して、TPUのバッチサイズは1024としました。詳しい条件はリポジトリを参照してください。

これはTPUが速すぎるというよりも、計算量が少なすぎて他のボトルネックのほうが支配的になっているから、と考えることができます。つまり、CNNだとGPUよりTPUのほうが速くなりやすいというよりも、**パラメーター（計算量）が物凄い多いCNNだとGPUよりTPUのほうがより相対的に速くなる**というほうが適切であるかのように思われます。

CNNのパラメーター数と計算量はおおよそ比例するので、**要はパラメーターの多いCNNを作ればTPUの本気が見れる**ということにもなります。ただし、ここには**1つ失敗例があってあまりに深い（パラメーター数84.9M、レイヤー数1000以上）NASNet Largeを使って学習させたところ、TPUの重みの初期設定で失敗してハングアップする**という現象がありました。ちなみにNASNet LargeはGPUでバッチサイズ1024にしても訓練できます[^2]。詳しくは[こちらに書きました](https://blog.shikoan.com/tpu-nasnet-large-failed/)。どこまでいけるのかは試していませんが、あまりにも深いモデルはダメということでしょうか。

[^2]: 訓練させたデータはCIFAR-100です。今回も同じデータを使います。

## パラメーター1億オーバーのWideResNet
パラメーターが多くする必要があるので、今回[WideResNet](https://arxiv.org/abs/1605.07146)を改造して作りました。ちなみにオリジナルのWideResNetのCIFARの構成は以下のとおりです。

![tpu_resnet_1.png](https://qiita-image-store.s3.amazonaws.com/0/232088/52964452-4f7a-b3d8-3ce9-f600d20c1b15.png)

ここで**kはチャンネルの大きさ（モデルの厚み）、Nは層の深さ（モデルの深さ）**を表します。今回は極端に深くしたくないので、**k=7, N=4**としました。そして今回の大きな変更点は、**オリジナルでは3x3の畳み込みを7x7の畳み込みに変更した**ことです。これでパラメーター数が49÷9≒5.4倍ぐらいになります。あくまで実験用にパラメーター数を増やすためのテクニックなので、通常はカーネルサイズを7x7にする必要はまずないです。

モデル生成のKerasのコードは以下の通りです。

```python
import tensorflow as tf
from tensorflow.contrib.tpu.python.tpu import keras_support
from tensorflow.keras.layers import Input, Conv2D, BatchNormalization, Activation, Add, AveragePooling2D, GlobalAveragePooling2D, Dense
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.models import Model

def create_residual_blocks(input_tensor, base_ch, k, N):
    start_tensor = input_tensor
    for i in range(N):
        x = Conv2D(base_ch*k, 7, padding="same")(start_tensor)
        x = BatchNormalization()(x)
        x = Activation("relu")(x)
        x = Conv2D(base_ch*k, 7, padding="same")(x)
        x = BatchNormalization()(x)
        x = Activation("relu")(x)
        x = Add()([start_tensor, x])
        start_tensor = x
    return x

# WideResNet
def create_wideresnet(k, N, use_tpu):
    input = Input(shape=(32, 32, 3))
    # conv1 : 32x32
    x = Conv2D(16*k, 1)(input)
    x = create_residual_blocks(x, 16, k, N)
    # downsampling 32->16
    x = AveragePooling2D(2)(x)
    x = Conv2D(32*k, 1)(x)
    # conv2 : 16x16
    x = create_residual_blocks(x, 32, k, N)
    # downsampling 16->8
    x = AveragePooling2D(2)(x)
    x = Conv2D(64*k, 1)(x)
    # conv4 : 8x8
    x = create_residual_blocks(x, 64, k, N)
    x = GlobalAveragePooling2D()(x)
    x = Dense(100, activation="softmax")(x)

    model = Model(input, x)
    model.compile(Adam(), loss="categorical_crossentropy", metrics=["acc"])

    if use_tpu:
        tpu_grpc_url = "grpc://"+os.environ["COLAB_TPU_ADDR"]
        tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(tpu_grpc_url)
        strategy = keras_support.TPUDistributionStrategy(tpu_cluster_resolver)
        model = tf.contrib.tpu.keras_to_tpu_model(model, strategy=strategy)
    return model
```

TPUとGPU両方対応させるために、tensorflow.kerasで定義しました。summaryを見てみます。長いので全体は[こちら](https://gist.github.com/koshian2/ae7bdbcfdbd93800bff52cd3107ce6f0)。**総パラメーター数は103Mと1億超え**です。

```:WideResNet
__________________________________________________________________________________________________
Layer (type)                    Output Shape         Param #     Connected to                     
==================================================================================================
input_1 (InputLayer)            (None, 32, 32, 3)    0                                            
__________________________________________________________________________________________________
#####
# 中略
#####
global_average_pooling2d (Globa (None, 448)          0           add_11[0][0]                     
__________________________________________________________________________________________________
dense (Dense)                   (None, 100)          44900       global_average_pooling2d[0][0]   
==================================================================================================
Total params: 103,465,028
Trainable params: 103,452,484
Non-trainable params: 12,544
```

これを使ってColabのGPUとTPU両方で訓練時間を比較していきます。

## 実験設定
* CIFAR-100をパラメーター103MのWideResNetで分類
* 変える条件は**デバイス（GPU,TPU）、バッチサイズ、Validationデータを使うか、DataAugmentationを使うか、データをどこから読ませるか（メモリorストレージ）、並列化するか（worker=1,4）**
* DataAugmentationはImageDataGeneratorでできる、ランダムクロップ＋水平反転という基本的なもの
* ストレージから読み込む場合は、一度バイナリのCIFARを1枚1枚のpngファイルに書き出し（[詳しい方法はこちら参照](https://blog.shikoan.com/cifar-to-image-file/)）、それをImageDataGeneratorのflow_from_directoryで読ませる。DataAugmentationなしでストレージから読み込む場合は、0.0～1.0のスケール変更のみ入れる。
* 20エポック訓練させて、2エポック目以降のエポックあたりの時間の平均を取る
* それ以外のパラメーターはだいたいデフォルト

デバイスごとにバッチサイズを以下の値で試します。

* **GPU:64, 128, 256**
* **TPU: 64, 128, 256, 1024, 2048, 4096**

これらはデバイスのメモリ制約によるものです。GPUに512やTPUに8192がないのはデバイスのメモリが溢れてしまうためです。

また、デバイス、バッチサイズ別に以下の8パターンを試します。また、各パターンについて、**(Validationデータを使うか, DataAugmentationを使うか, データをストレージから読ませるか)**の順でTuple形式で条件を略して表記します。

| パターン | Val | Aug | ストレージ |    表記   | 並列 |
|:--------:|:---:|:---:|:----------:|:---------:|:----:|
|     1    |  ☓  |  ☓  |      ☓     | (F, F, F) |   1  |
|     2    |  ○  |  ☓  |      ☓     | (T, F, F) |   1  |
|     3    |  ○  |  ○  |      ☓     | (T, T, F) |   1  |
|     4    |  ☓  |  ☓  |      ○     | (F, F, T) |   1  |
|     5    |  ○  |  ☓  |      ○     | (T, F, T) |   1  |
|     6    |  ○  |  ○  |      ○     | (T, T, T) |   1  |
|     7    |  ☓  |  ☓  |      ○     | (F, F, T) |   4  |
|     8    |  ○  |  ○  |      ○     | (T, T, T) |   4  |

この8パターンを、デバイスバッチサイズごと（9パターン）に計72ケース測定します。

コードはこちらにあります。https://github.com/koshian2/TPU-Benchmark

## サマリー
8パターンのうち最もデバイス、バッチサイズ間の差が大きかったパターン3（T, T, F）を示します。全体の結果は[こちら](https://github.com/koshian2/TPU-Benchmark/blob/master/result/result.md)

| デバイス | バッチサイズ | 訓練前 |   ep.1 | ep.2- 平均 | ep.2- 標準偏差 | 20eps（実測） |
|:--------:|:------------:|-------:|-------:|-----------:|---------------:|:-------------:|
|    gpu   |      64      |    0.5 | 1130.8 |     1108.4 |            0.9 |    6:09:52    |
|    gpu   |      128     |    0.6 |  852.2 |      821.3 |            2.6 |    4:34:17    |
|    gpu   |      256     |    0.5 |  769.7 |      742.9 |            1.3 |    4:08:05    |
|    tpu   |      64      |    0.8 |  907.5 |      844.3 |            6.9 |    4:42:30    |
|    tpu   |      128     |    0.6 |  507.0 |      421.7 |            2.6 |    2:22:00    |
|    tpu   |      256     |    0.9 |  394.3 |      234.6 |            1.6 |    1:20:52    |
|    tpu   |      512     |    0.4 |  297.5 |      136.9 |            0.8 |    0:48:19    |
|    tpu   |     1024     |    0.4 |  247.2 |       80.3 |            0.2 |    0:29:34    |
|    tpu   |     2048     |    0.8 |  162.2 |       58.1 |            0.8 |    0:21:06    |
|    tpu   |     4096     |    0.7 |  145.5 |       45.6 |            1.7 |    0:16:52    |

実測値ベースでは、**同一バッチサイズで比較するとTPU-256はGPU-256の3.07倍速く、異なるバッチサイズでの比較を許容するとTPU-4096はGPU-64の21.93倍速い**結果となりました。

また、GPU・TPUとも1エポック目のみ遅く（TPUは特に顕著）、2エポック目以降はより短い時間で安定的に訓練が進むという傾向があるので、エポック数が多くなれば、全体の訓練時間の比は2エポック目以降のエポックあたりの時間の比に収束します。この「ep.2～平均」で比較すると、**同一バッチサイズならTPU-256はGPU-256の3.17倍速く、異なるバッチサイズならTPU-4096はGPU-64の24.33倍速い**という結果になりました。

ただ、異なるバッチサイズでの比較を許容するのなら、TPU-64はGPU-256の1.14倍遅いというのも言えるため、**必ずしもTPUが対GPUで最速であるということを保証できるものではありません。**しかし、トップ差でTPUはGPUの20倍以上速くなるという結果も同様に出ているので、その点はしっかり受け止めるべきではないかと思います。

この時間の差を体感するために、エポック数を変えて訓練した場合の総時間を推定してみました。計算式は「[訓練前]＋[ep.1]＋[ep.2平均]×([エポック数]-1)」です。実践的に(T, T, T)で並列数1＝パターン6での一覧です。

| デバイス | batch |  10 eps | 20（実測） |   50 eps |  100 eps |  200 eps |   500 eps |  1000 eps |  2000 eps |
|:--------:|:-----:|--------:|-----------:|---------:|---------:|---------:|----------:|----------:|----------:|
|    gpu   |   64  | 3:10:20 |    6:20:15 | 15:50:02 | 31:39:40 | 63:18:55 | 158:16:41 | 316:32:58 | 633:05:32 |
|    gpu   |  128  | 2:22:04 |    4:43:50 | 11:49:11 | 23:38:04 | 47:15:52 | 118:09:15 | 236:18:14 | 472:36:11 |
|    gpu   |  256  | 2:09:58 |    4:19:20 | 10:47:29 | 21:34:23 | 43:08:11 | 107:49:34 | 215:38:34 | 431:16:32 |
|    tpu   |   64  | 2:14:30 |    4:27:24 | 11:06:07 | 22:10:39 | 44:19:43 | 110:46:54 | 221:32:13 | 443:02:51 |
|    tpu   |  128  | 1:10:11 |    2:18:33 |  5:43:39 | 11:25:29 | 22:49:10 |  57:00:12 | 113:58:35 | 227:55:21 |
|    tpu   |  256  | 0:47:12 |    1:27:15 |  3:27:23 |  6:47:35 | 13:28:01 |  33:29:19 |  66:51:29 | 133:35:48 |
|    tpu   |  512  | 0:30:28 |    0:53:25 |  2:02:16 |  3:57:01 |  7:46:31 |  19:15:03 |  38:22:34 |  76:37:38 |
|    tpu   |  1024 | 0:20:55 |    0:35:02 |  1:17:23 |  2:27:57 |  4:49:07 |  11:52:36 |  23:38:24 |  47:10:00 |
|    tpu   |  2048 | 0:13:20 |    0:24:30 |  0:57:59 |  1:53:47 |  3:45:22 |   9:20:10 |  18:38:09 |  37:14:07 |
|    tpu   |  4096 | 0:12:21 |    0:22:17 |  0:52:06 |  1:41:47 |  3:21:09 |   8:19:15 |  16:36:06 |  33:09:47 |

GPU-64というのが極端な値だとしてもバッチサイズを大きくできれば、仮に2000エポックも訓練させる場合は、**18～25日（430時間～630時間）vs2日前後**という大きな差が生まれます。正直2000エポックも訓練させるのにColabを使うのは本来の使い方ではありませんが、クラウド版TPUもあるので[^3]TPUで長時間訓練させることが不可能ではありません。

[^3]: Colab版はTPUv2、有料のクラウド版はTPUv3なので、おそらくColab版よりももっと速いはずです。

機械学習自体が反復のプロセスなので、1つのモデルを訓練するのに1ヶ月近くかかるのと2日で終わるのでは、時間あたりの反復回数が全然変わってきます。それはつまりより多くのアイディアを試せることになるので、結果的には時間あたりの生産性を増やしたり、精度を向上できる可能性が大きくなることを意味します。**「速さは正義」**なのです。

また、最近（2018年9月）にState of the Artを達成した[BigGAN](https://arxiv.org/abs/1809.11096)というモデルによれば、**バッチサイズを大きくすることがGANの安定化に寄与する**という研究があります。これはGANの場合ですが、バッチサイズを大きくすることによって訓練速度以外のメリットを享受できるという明確な証拠でもあります。一般的にはバッチサイズを変えることはハイパーパラメータをチューニングし直す必要があり、面倒に感じるかもしれませんが、これから新規でモデルを作る場合は、もしメモリの要件が許すのならバッチサイズを大きくして作ったほうがいいのかもしれません。ちなみにこの**BigGANはTPUで訓練している**とのことです。

次に細かく条件を分けて見ていきます。

## 詳細な結果
### 1.バッチサイズ別
バッチサイズ別の2エポック目以降のエポックあたりの平均秒をグラフで示したものです。(F, F, F)のパターン1で比較しました。バッチサイズを大きくするとTPUは非常に速くなるのは確認の通りです。
![01_batch_size.png](https://qiita-image-store.s3.amazonaws.com/0/232088/7af9662a-50a2-2179-50d4-ae9d4a37e248.png)
グラフの縦軸を対数スケールにしてみました。
![02_batch_size_log.png](https://qiita-image-store.s3.amazonaws.com/0/232088/20ee936c-9b58-ac08-deee-843a18181105.png)
これを見ると、**TPUは対数スケールではほぼ直線的に速くなっている**のがわかります。GPUは3ケースなのでなんとも言えませんが、256あたりで頭打ち感が否めません。TPUも4096あたりは若干頭打ちな印象はありますね。

### 2.Validationの有無
次にエポックの最後にValidationデータを入れることで訓練時間がどの程度変わるかを見ます。CIFAR-100のテストデータをValidationとして入れたため、Validationの有無でデータ数が5万と6万の違いがあります。これは(F, F, F)のパターン1と、(T, F, F)のパターン2を比較したものですが、データ数が1万増えているので、Validationを入れたほうが遅くなるのは当たり前と言えるでしょう。
![03_use_validation.png](https://qiita-image-store.s3.amazonaws.com/0/232088/2778ee64-2056-c858-803e-80478942e989.png)
しかし、データ数が1.2倍の割にはそこまで処理時間は増えていないように見えます。(T, F, F)の時間を(F, F, F)の時間で割って、Validationの有無による時間比を求めてみます。
![04_use_validation_ratio.png](https://qiita-image-store.s3.amazonaws.com/0/232088/fc5416eb-6908-0364-5f6a-30c284a6c213.png)
Validationを入れたほうが時間が増えるのはそうだとしても、データ数が1.2倍になっているのに処理時間は1.2倍までなっていませんね。おそらくですが**ValidationだとBack Propagationを計算しないため速い**[^4]のでしょう。Validationのデータ数は訓練データよりもずっと少ないことが多いので、これだったら**GPU/TPU問わずどんどんValidationを入れたほうがいい**ように思えます。モデルのチェックポイント作ったり何かと有用なので。

[^4]: Back Propは勾配降下法で用いる偏微分を計算するためのものです。Validationでは誤差（Fore Prop）のみ計算できればよくてパラメーターを更新しないため、Back Propを計算する必要がありません。

### 3.DataAugmentationの有無
(T, F, F)のパターン2と(T, T, F)のパターン3を比較しました。予想に反してDataAugmentationをしたほうが速くなるという結果が出ました。
![05_use_augmentation.png](https://qiita-image-store.s3.amazonaws.com/0/232088/3f509a09-e19f-ffc1-e84e-c7339756b75c.png)
DataAugmentationをしたパターン（T, T, F）をしなかったパターン（T, F, F）で割ったものが以下のとおりです。DataAugmentationをしたほうが最大で1割近く速くなっているという奇妙な結果となりました。
![06_use_augmentation_ratio.png](https://qiita-image-store.s3.amazonaws.com/0/232088/e4d6f30c-c0b2-24bb-5d20-6708b44a3d13.png)
これには2つ理由が考えられると思います。

* DataAugmentationを用いる際はImageDataGeneratorを使い、model.fitではなくmodel.fit_generatorを使っているため**DataAugmentationが速いのではなくfit_generatorが速い(fitが遅い)説**
* **fit_generatorではバッチサイズの端数**（例えばバッチサイズが4096だったら、4096×12=49152個のサンプルを使い、残りの848個のサンプル）を**捨てているため速い説**

この2点については後ほど検証します。DataAugmentationは明らかに余計な処理であるため、これ自体が速度の向上に寄与しているとはまず考えにくいです。ただ、DataAugmentationの操作にもよりますが、ランダムクロップ＋水平反転という基本的な操作ならば、**DataAugmentationはTPUの速度すら損ねない（ほとんどボトルネックにはならない）レベル**であるというのは言うことはできそうです。

### 4.ストレージから読み込むかどうか
常識的には「メモリから読み込むと速い、ストレージ（HHD,SSD）から読み込むと遅い」ということは考えられますが、ストレージからデータを読み込ませた場合どの程度速度の変化が現れるのでしょうか。(F, F, F)のパターン1と(F, F, T)のパターン4を比較してみます。
![07_from_storage.png](https://qiita-image-store.s3.amazonaws.com/0/232088/6364fe40-6fab-3c27-0e0c-45c480054277.png)
これも意外な結果になって、「ストレージから読み込ませると速い」という結果になりました。
![08_from_storage_ratio.png](https://qiita-image-store.s3.amazonaws.com/0/232088/d1a9985b-4587-cdf3-cc52-91ad6ee2a46a.png)
ただ、よく考えてみると、これもDataAugmentationの場合とほぼ同じで、**結局はfitとfit_generatorの差**だよなという印象は拭えません。ストレージから読み込む場合は、ImageDataGeneratorのflow_from_directoryを用いました。

本来はUnixBenchを動かしてストレージの読み込み速度を見るべきですが、上手く動かなかったので諦めました。一般的にはストレージの読み込み速度を見ると、シーケンシャルリードがHDDで150MB/s、SSDで500MB/s（あくまで大雑把な目安で物によって全然違う）ぐらいなので、おそらくCIFAR程度だとストレージからの読み込みがボトルネックにならないのかと思います。CIFAR-100の6万枚をpngファイルに書き出してtar.gzで固めたところ126MBでしたので、1周50秒で訓練させたとしてもこのケースではまだまだストレージ側は余裕があるはずです。

### 5.並列化をするとどうか
自然に考えると「並列化したほうが速い」と思われますがどうでしょうか。並列化をする場合は、model.fit/fit_generatorの「workers=4, use_multiprocessing=True」としました。(T, T, T)で並列数1（パターン6）と並列数4（パターン8）を比較しました。

![09_num_workers.png](https://qiita-image-store.s3.amazonaws.com/0/232088/ffdadc24-15eb-c549-6b5b-7f65b30faa4d.png)

GPU-256の並列数1がすごく遅くなっていますが、多分環境の調子が悪かったのだと思います。これだけ2エポック目以降の時間の標準偏差が他の50倍近くあるので。
![10_num_workers_ratio.png](https://qiita-image-store.s3.amazonaws.com/0/232088/ff09d1dd-70ea-9ca8-ea8e-74570062b51a.png)
これは並列数4の時間を並列数1で割ったものです。この数字が小さくなるほど並列化の効果が出ることになります。GPUやTPUの少ないバッチ数（～1024）ではある程度は並列化の効果はあるものの、**TPUでバッチサイズが大きくなったときに並列化すると逆に遅くなる**というのが確認できました。

## 追試：fit vs fit_generator
ここで1つ大きな疑問がわきます。DataAugmentationとストレージのところであった、「実はこれらが高速化に寄与しているのではなくてfit_generatorが速いからでは」という疑問です。そこで次の3つのモードを定義して追試することにしました。

1. 訓練データ数5万のままfit
2. 訓練データ数をバッチサイズの倍数になるように切り詰めてfit（例：バッチサイズ4096なら12×4096=49152とする）
3. fit_generatorを使う

上からモード1、モード2、モード3とします。generatorは0.0～1.0のスケール調整だけ行ったもので、データは全てメモリから読み込ませています。それ以外の条件は前の実験と同じです。TPUでバッチサイズを256、512、1024、2048、4096と変えて試してみます。結果は以下の通りです。

| batch/mode |   1   |   2   |   3   |   2÷1 |   3÷1 |
|:----------:|:-----:|:-----:|:-----:|------:|------:|
|     256    | 219.9 | 221.0 | 219.8 | 1.005 | 0.999 |
|     512    | 131.3 | 131.6 | 128.3 | 1.002 | 0.978 |
|    1024    |  80.1 |  81.6 |  79.7 | 1.018 | 0.995 |
|    2048    |  61.2 |  60.1 |  58.2 | 0.982 | 0.951 |
|    4096    |  50.9 |  48.7 |  46.0 | 0.956 | 0.903 |

1,2,3のところは2エポック目以降の平均時間を記したもの、2÷1はモード2の時間をモード1で割ったもの、3÷1はモード3をモード1で割ったものです。

![11_fit_fitgenerator.png](https://qiita-image-store.s3.amazonaws.com/0/232088/c05474ca-ea58-5a04-11a0-e1c8fab4f909.png)

これらを見ると、**先程の仮説はどちらも正しい**ということが言えそうです。つまり、

* **fitをそのまま使うよりも、バッチサイズの倍数になるようにサンプル数を切り詰めてfitを使ったほうが若干速い**
* **fit_generatorを使うともっと速くなる**

です。なので、**DataAugmentationやストレージの有無の差は、どちらかというとfit vs fit_generatorの差や訓練データ数がバッチサイズの倍数かどうかの差で説明できる、むしろそう考えるほうが自然**といえるでしょう。同時に、**DataAugmentationや読み込み元の差は、今回の場合は特にボトルネックとなることはなかった、つまりこれらの要素を高速化のために除外する必要はない**ということになります。

## まとめ
とても長くなってしまいましたが、結論をまとめましょう

* 1億パラメーターのCNNでColabのTPUvsGPUを比較したところ、異なるバッチサイズ間の比較を許容すれば**TPUはGPUより20倍以上速くなる（24.33倍）**
* 機械学習は反復のプロセスなので「速さは正義」
* **速度を最も決定づけるのがバッチサイズ**で、バッチサイズを大きくできればTPUは指数関数的に計算性能が強くなる（時間ベースでは対数関数的に減少）。またBigGANではバッチサイズを大きくすることで訓練を安定化しているので、バッチサイズを大きくすることには速度面以外のメリットもあることが確認されている。
* Validationはデータ数が増えるため若干遅く（5万：1万で数％）はなるが、Back Propがないためそこまで遅くならない。入れるのをためらう必要はない。
* model.fitよりも、データ数をバッチサイズの倍数になるようにしてfitさせる、fit_generatorを使うほうが若干高速化する
* DataAugmentationやデータの読み込み元（メモリorストレージ）の条件は、今回の場合TPUのボトルネックとなることはなかった
* 並列化はバッチサイズが大きい場合は逆効果になることもある

1つ補足ですが、今回やった実験はかなり特殊な条件なので、今回24倍速くなったからといって必ずしもこの速度が出る、もっと言えばTPUはGPUより速くなることを保証できるものではありません。現に前回の投稿で確認したように、MLPではGPUのほうがまだ速いです。ただ、**ちょうどハイパーパラメータの選択のように、デバイスをGPUとTPU両方変えて訓練してみる、そして速度だけではなく総合的に見てプロジェクトにあわせて良いほうを選ぶというのが重要ではないか**と思います。

最後にKerasの作者のFrançois Cholletさんの言葉で締めくくりたいと思います。

<blockquote class="twitter-tweet" data-lang="ja"><p lang="en" dir="ltr">You can start training Keras models on TPUs, from the comfort of your browser, in a few seconds. Try it: <a href="https://t.co/G95rxp7U7P">https://t.co/G95rxp7U7P</a><br><br>I expect that most industry-scale deep learning workflows will move from GPUs to TPUs over the next couple of years. <a href="https://t.co/WG63DHbuzD">pic.twitter.com/WG63DHbuzD</a></p>&mdash; François Chollet (@fchollet) <a href="https://twitter.com/fchollet/status/1047841218539925504?ref_src=twsrc%5Etfw">2018年10月4日</a></blockquote>
<script async src="https://platform.twitter.com/widgets.js" charset="utf-8"></script>

「私は、業界規模の大規模な学習ワークフローが今後数年間でGPUからTPUに移行することを期待しています。（Google翻訳による）」とのことです。