TPUもリリースされ、CPU, GPU, TPUをどういった時に、何を基準に使い分ければよいのかと思っていたら、以下に記載されていることだったのでメモ。
https://cloud.google.com/tpu/docs/tpus
Cloud TPUは特定のワークロードに最適化されているので、機械学習のワークロードによってはCompute EngineのインスタンスでCPUやGPUを使用した方がよいケースがある。一般的には、どの環境を選択すべきかは以下のガイドラインに従えばよい。
CPU
- 柔軟性があるプロトタイプを素早く作る必要がある場合
- モデルのトレーニングに時間がかからないような単純なモデル
- 小さなモデル
- C++で記述されたカスタムTensorFlow operationが大部分を占めるモデル
- 利用可能なI/Oやホストのネットワークバンド幅が制限されているモデル
GPU
- TensorFlowで記述されてないモデルやTensorFlowで記述できないモデル
- ソースが存在していなかったり、変更するのがあまりにも煩わしいモデル
- CPUで一部実行しないといけないようなカスタムTensorFlow operationがある膨大にあるモデル
- Cloud TPUで利用できないTensorFlowオペレーションがあるモデル
- 大きなバッチサイズを用いる中規模から大規模なモデル
TPU
- 行列計算が大きく占めているモデル
- 主要なトレーニングループの中にカスタムTensorFlow operationがないモデル
- トレーニングに数週間や数ヶ月かかるモデル
- 大きなバッチサイズを用いる大規模なモデル
Cloud TPUに向いてないワークロード
- 頻繁にブランチングしてたり要素ごとの積が多くを占めるような線形代数のプログラム。TPUは大量の行列積演算が速くなるように最適化されている。そのため行列積が多くを占めてないワークロードはTPU上では十分なパフォーマンスがでない。
- TPU上で利用できないようなスパースにメモリアクセスするワークロード
- 高精度演算を要求するワークロード。例えば倍精度はTPUに向いてない
- CD++で記述されたカスタムTensorFlow operationを含むニューラルネットワークのワークロード。特にメインのトレーニングループの中にカスタムオペレーションが含まれているケースはTPUに向いてない
参考） TPUで利用可能なTensorFlow ops
https://cloud.google.com/tpu/docs/tensorflow-ops