TPUを仕組み理解からTensorFlowで実装まで

TPUって何?

見ていると気になることはいっぱいある
上げればキリがないものだが、昔のメモをさグッたらTPUという文字が出てきたのでこの際に調べておこう。

Tensor Processing Unit

最も事細かにかかなので、あしからず。
仕組みの説明はGoogleCloudの説明がしっくりきた。
まず、CPU・GPUの復習

  • CPU
    • ALUに何でもやらせるから柔軟性は高い。
    • でもキャッシュ読み込み→命令処理→書き込みまでが時間とエネルギーの無駄
    • パイプライン処理で頑張っても、結局たかが知れてる
  • GPU
    • 1個のプロセッサに大量のALUを入れたもの。単純な命令を「数」で倒せる。
    • 数の分だけ高速化するが、電力を食う。
    • CPUと仕組みは同じなので、中間でキャッシュを経由するため時間とエネルギーの無駄
      • この呪縛をノイマンボトルネックというらしい。この名は初めて知った

Googleはうまいことを考えた。GPU・CPUの欠点を克服するには、構造をかえようと
キャッシュに行く前にどんどん足してしまえばよい!
Google公式のGIFがよいね

データを左から列方向に展開して→に進めていく

パラメータを下から行方向に展開して↑に進めていく

NNっていうのは入力データに対し、パラメータを1対1(y=ax+b)でつなげるものなので、
上2図のように流していけば、否が応でも演算器にぶち当たり計算せざるを得ない。
出口である右端では最終的に下みたいな行列演算(matmul)ができていることになる

感動した。頭いい。

このキャッシュ通さない作戦をシストリックアレイって言うらしい。
また、NNの計算過程では精度はそこまで重要ではないらしく、
量子化を行って一般的な32/64bit演算から8bit演算の仕組みを採用したため、実質4倍高速化(=省エネ)
この精度周りの下りは結構アツいらしい
https://ascii.jp/elem/000/004/014/4014066/

ムーアの法則って呪縛に聞こえてきた…

GPUとCPUのお互いいいところがあるように、TPUにも向き不向きがある。
Googleから直接拝借する。

・CPU
 最大限の柔軟性を必要とする迅速なプロトタイピング
 トレーニングに時間がかからない単純なモデル
 実際のバッチサイズが小さい小規模なモデル
 C++ で記述されたカスタム TensorFlow 演算が多くを占めるモデル
 ホストシステムの使用可能な I/O またはネットワーク帯域幅によって制限が課せられるモデル
・GPU
 ソースが存在しないモデルまたはソースを変更するのが煩雑すぎるモデル
 CPU 上で少なくとも部分的に実行しなければならない多数のカスタム TensorFlow 演算を使用するモデル
 Cloud TPU で利用できない TensorFlow 演算を使用するモデル(利用可能な TensorFlow 演算のリストをご覧ください)
 実際のバッチサイズが大きい中〜大規模なモデル
・TPU
 行列計算が多くを占めるモデル
 メインのトレーニング ループ内にカスタム TensorFlow 演算がないモデル
 トレーニングに数週間または数か月かかるモデル
 実際のバッチサイズが非常に大きい非常に大規模なモデル

https://cloud.google.com/tpu/docs/tpus?hl=ja

GPUはある程度複雑な計算でもある程度応用(適応?)が効いたが、TPUは一切それが効かない。
その代わり、単純な行列演算の塊なら必殺!といった次第だろう

ColabとかKaggleとかでは使えるのでつかってみた

実装してみます。とは言っても他人コードを拝借する。一部加えつつ、、
1年前にkaggaleでTPUつかおうぜ!みたいなのを思い出したので、それを題材にします。
お花分類ですね

こちらに書いてあるTPUの使い方はとてもわかりやすいと思います(あくまで理解のうえでは)
SAMPLEコードにめもを追加していきます。

TPUとGPUを切り替える部分。
各TPUを制御するcluster_resolverを使ってTPUをネットワークから探索。
各TPU間の同期方法を定義するstrategyは必須。strategyはあくまでインスタンスなので、のちほど定義する
TPUStrategyの場合、アクセラレーターで分散する際に変数のミラーリングがされ、分散学習の重みを集約→重みを更新(同期)→共有する。この仕組みをall-reduceっていう。
TPUは使えない場合はGPUの同じミラーリング戦略であるMirroredStrategyをつかう
以下、日経 xTECHを拝借

まずはデータを撮ってくる。
バッチサイズが大事で、ここは8の倍数(GoogleTPUv3.8の場合8bit量子化しているので)にしないと行けない。
これは後で使うSteps_per_executionも同じことが言える。

学習率を途中で変えるらしいです。TPUに限定してなぜカスタマイズするのかは不明。
ただ自分自身で変更するのは初めて。よい知見。

この辺は画像を撮ってくるだけ。TFRecordなので、それに応じて読み込む。
ここは全て借り物。おパクり。
重要な点は最初に設定した”AUTO”。
TPU・GPUはPrefetchが速度のボトルネックなのでここは大事。それ以外はデータよりけりだと思う

よしなに分割。

モデル作成。

strategy.scope()でくくることで、先程のインスタンスがようやく登場。
この中の処理のみがTPUで共有される。つまり、モデル部分だけが共有される。
今回の例ではSAMPLEに従って、xceptionを採用。下記の通り他のモデルを使いたいならテキトウに変更
steps_per_execution=8は、1callでなんバッチ読み込むかを定義する。上の”AUTO”と同じで高速化に貢献
単純なギモンとしてGPUでも同じこと言えるよね?と思ったが、GPUは関係ないらしい。

学習。もはやGPUのときと変わらない。

このあと保存と読み込みがあるが、GPUと読み込み方法は異なるものの、余り本質ではないんので省略。
いまはTenswflow liteというものがあって、いわゆるエッジコンピューティング用に特化されたフレームワークがあるけど、そこでも活用できるみたい。
電子工作でおなじみのswitchScienceから出ている。今度遊んでみたい。1万するけどwww
https://www.switch-science.com/catalog/5817/

gmoがなんかやっていた。

https://recruit.gmo.jp/engineer/jisedai/blog/edge_tpu_tf_lite_basics/

感想

  • TPUナニそれすごそう!という印象だったが、案外GPUと実装の手間がほぼ変わらない。クラウドならかんたんに使えそう
  • 久しぶりにCSチックな言葉が出てきたので興奮した。大学院のときよりも

補足