Googleの機械学習用チップTPUの最新バージョンTPU v3 Podを試してみた

Tensor Processing Unit（TPU）とは

近頃は機械学習やニューラルネットワークによる予測の精度はより正確になってきています。理由の１つとして、大量のパラメータを持つ複雑なモデルをたくさんのデータでトレーニング（学習）することが可能になったことが挙げられるでしょう。
しかし、こういったモデルをトレーニングする計算時間とマシン費用に対する問題は依然として残っています。大量のデータとパラメータを短時間で計算するには強力なコンピューティングリソースが必要となるので当然費用がかかります。かと言って、貧弱なマシンで計算すると計算時間が数日、数週間かかってしまい要件を満たせなくなってしまう場合があります。従って、安くて強力なプロセッサが常に求められています。
そこでGoogleはTPU（Tensor Processing Unit）を開発しました。TPUは機械学習用に設計されたASIC（特定用途向け集積回路）です。
第1世代のTPU v1は2016年に発表され、最新のTPU v3は2018年に発表されました。
TPUは行列の演算を効率よく行うことに特化しています。これはニューラルネットワークのように同じような行列計算を何度も繰り返すモデルに特に効果を発揮してくれます。

では、普通のPCに内蔵されているCore Processing Unit (CPU)と何が違うのでしょうか。
Google Cloud Platform Japan ブログでわかりやすい記事が紹介されていました。
この記事ではニューラルネットワークを用いて画像から数字を判別するモデルを挙げて説明しています。
例えば”8”を識別する場合、モデル内部ではイメージとしては以下のようなことを行っています。

この演算はCPUで行うときは以下のように動作しています。

CPUは演算がひとつずつ行われるイメージです。

一方でTPUは計算ごとにメモリへの書き込みを行わずに、バケツリレーのように一気に演算が行われます。演算中は一度もメモリへの書き込みが発生しないために、処理速度のボトルネックが解消されてます。行列演算では内部の加算・乗算を並列計算可能であるからこそ、このような処理が実現できます。


また、一般的なCPUでは32-bitまたは16-bitの浮動小数点での計算処理を行っているのに対して、TPUでは量子化を行うことにより8-bitの計算を行っています。演算の精度を犠牲にする代わりにメモリ消費を抑え、消費電力を減らすこともできます。しかも精度を犠牲にしているといっても、工夫を少し行うことで十分な精度のモデルを作成することが可能となっています（詳しくはこちら）。

GCPのプロダクトとしてのCloud TPU

GCPではこのすさまじい処理能力を持つTPUを利用することができます。
また、TPU Podと呼ばれる数百個のTPUのチップと数十台のホストマシンで構築された密結合型の高スペックプロセッサも利用することができます。単体のTPUが1チップ(2コア)に対して、v2 Podのフル構成は256チップ(512コア)、v3 Podのフル構成は1024チップ(2048コア)となっています（参照：TPUチップ）。
以下にTPU v2とTPU v3、TPU v2 Podフル構成、TPU v3 Podフル構成の性能表を示します。

※FLOPS（フロップス、Floating-point Operations Per Second)
コンピュータが1秒間に処理可能な浮動小数点演算の回数を示す単位。
※HBM（High Bandwidth Memory）

Graphics Processing Unit（GPU）の代表格としてNVIDIA Tesla V100がありますが、単精度で14 TFLOPS(テラFLOPS)だそうです。（参照：Tesla V100）
Cloud TPU v3はこれの420÷14=30倍の性能です。TPU v3 Podフル構成に至っては100,000÷14=約7千倍という性能となっています。そもそも100PFLOPS(ペタFLOPS)はスパコンの性能ランキングで4位の天河二号の理論性能値（Rpeak）と同等かそれ以上という化け物っぷりです。(参照：TOP500 LIST)

2019/7/19現在ではTPUとTPU Podを利用可能な地域は限られており、リージョンと利用できるTPUの組み合わせは次のようになります。

Cloud TPU v2 : US、ヨーロッパ、アジア
Cloud TPU v3 : US、ヨーロッパ
Cloud TPU v2 Pod（ベータ） : US, ヨーロッパ
Cloud TPU v3 Pod（ベータ） : ヨーロッパ

Podに関してはクオータ（割当）申請が必要なものもあります。ドキュメントとコンソール上の”IAMと管理”から”割り当て”を確認しましょう。

実は意外と安いCloud TPU

これほどの処理能力を持つTPUですが、GCP上でのCloud TPUは意外と安いです。

執筆時点(2019/8/28)での料金は以下のようになっています。

※v3-32 Cloud TPU v3 Podのみeurope-west4、他はus-central1

シングルのTPU利用時にはプリエンプティブを指定することで安くすることもできます。（プリエンプティブの指定のやり方は次節参照）
一方でTPU Podではプリエンプティブを指定することができません。
上表からは1時間単位の課金体系のように見えますが、実際は1分単位の課金となります。またPodでは継続利用割引などもあるので公式ドキュメントを参考にすると良いでしょう。
v3-32より大きなPodはセールスチームに問い合わせましょう。

前述したGPU(NVIDIA Tesla P100)は1時間1.46$ですが、その50倍近い性能のTPU v3を1時間2.4＄で借りられるとは恐るべし。
さらにモデルのトレーニング時間は短くなるので、コスパ最強です！！！

Cloud TPU Podを使ってみた

さて、ドキュメント通りの性能かつ費用ならば、エンジニアにとってこれほど嬉しいことはないでしょう。これからは実際にCloud TPUを使ってみたいと思います。
公式のチュートリアルを参考にしてResNet-50（以下、ResNet）のトレーニングを行いました。ResNetは2015年のImageNetのコンペで優勝した優秀なモデルです。ResNetは層が深く、トレーニング時間が長いことで有名です。
用いるデータはfake_imagenetという、Googleが用意している有名なImageNetのデータ約128万を模したフェイクイメージデータです。
パラメータとしてトレーニングバッチサイズは4096、エポック数は90で行います（注）。
今回はCloud TPU v3 PodとCloud TPU v2 Podの比較検証とします。両方ともコア数は32です。

（注）
バッチサイズとはモデルが一回のトレーニングに使うデータ数のことです。エポック数とは全データをトレーニングする回数です。例えば、全てのデータが100、バッチ数が20だとしたら、1エポックは100÷20 = 5バッチになります。

クオータの申請（必要な場合）

2019/07/19現在TPU Podの利用にはクオータ（割当）申請が必要となります。こちらを参考にして事前に申請しておきます。
今回はv3-32を申請し、europe-west4-aのTPU Podを割り当てられました。以下でPodのコア数32個が割り当てられていることを確認できます。

Cloud TPU利用の環境設定

Cloud TPUを用いるにはVMを立てる必要があります。ここでやり方としては大きく2種類あります。

• ctpuコマンドを用いたシンプルな環境構築
• VMとCloud TPUを別々に立てて接続

今回はctpuというコマンド（Cloud TPU Provisioning Utility）を使ってVM+Cloud TPUの環境を構築します（参照：ctpu）。

VMとTPUのセットアップ

最初にGCPコンソール画面の右上にあるアイコンをクリックし、Cloud Shellを立ち上げます。

立ち上がったコンソール画面から以下のコマンドを入力します。

ctpu up \
--tpu-size=v3-32 \
--name=tpu-v3-32-test \
--zone=europe-west4-a \
--disk-size-gb=500 \
--machine-type=n1-standard-8

シングルTPUの場合、–preemptibleフラグを指定することでプリエンプティブTPUを利用することができます。

europe-west4-aゾーンに、VMインスタンスとv3-32のTPU Podが立ち上がります。
VMとTPUの名前はnameフラグで指定した名前（今回はtpu-v3-32-test）で立ち上がります。

VMが立ち上がったと同時にVMへのログインも完了します。

GCSバケットの用意

次にトレーニングの結果などを保存するためのGoogle Cloud Storage（GCS）バケットを作成します。TPUと同一リージョンのRegionalバケットを用意する必要があります（今回はeurope-west4のバケットを用意しました）。
このバケットにはトレーニングの結果が保存されます。

トレーニングの実行

この手順で作成されたVMには、/usr/share/tpu/models/official/resnet/ 下にサンプルプログラムが最初から用意されています。これを用いて学習を行いたいと思います。

プログラムの中身について少し触れておくと、このサンプルプログラムではTensorflowのtf.estimatorという高レベルAPIを使って実装されています。このプログラムはestimatorをTPU用にするTPU estimatorを使って実装されています。Estimatorの一部のコードを変更するだけなので比較的簡単に既存のEstimatorのコードをTPU用に切り替えることができます。

まずは環境変数を設定します。

バケット名

export STORAGE_BUCKET=gs://{YOUR_BUCKET_NAME}

前節で用意したバケットを指定します。

pythonのパスの設定

export PYTHONPATH="$PYTHONPATH:/usr/share/tpu/models"

これはそのままです。

使用するTPUバージョンの指定

export ACCELERATOR_TYPE=v3-32

使用するTPUのバージョンを指定してください。
（シングルTPU v2ならv2-8、v3 Pod 1024コアならv3-1024）

プログラムのあるディレクトリへ移動

cd /usr/share/tpu/models/official/resnet/

これはそのままです。

最後にプログラムを実行します。

python resnet_main.py \
  --tpu=${TPU_NAME} \
  --data_dir=gs://cloud-tpu-test-datasets/fake_imagenet \
  --model_dir=${STORAGE_BUCKET}/resnet_$(date +%s) \
  --mode=train \
  --config_file=configs/cloud/${ACCELERATOR_TYPE}.yaml

このコマンドから推定できるかもしれませんが、トレーニング用の教師データはgs://cloud-tpu-test-datasets/fake_imagenetを直接参照しています。

プログラムが実行されます。

バッチごとに読み込まれ、トレーニングが進みます。

TPUの詳細ページです。CPU（ここではTPUの使用率が正しいですが）が安定して使用されていることがわかります。

以下はトレーニングが終了した状態です。

トレーニング結果としてはロスとステップ数が表示されます。
詳細なデータはTensorBoardを用いることでリアルタイムでトレーニングやTPUの状況をプロファイルすることも可能です（設定方法は今回は割愛します）。

またTPU Podではevaluationを使うことができません。evaluation時には別のVMを立てて、バケットに保存されているチェックポイントデータからevaluationする必要があります。

クリーンアップ

シングルTPUは一時停止が可能です。一方でTPU Podは一時停止することができないので、停止するためにはリソースを削除する必要があります。
TPU Podを削除しないでそのままにしておくと、立てていた分の請求が来ます。安くない額が余計にかかってしますので、使い終わったら必ず削除しましょう。

VMからログアウト

exit

TPUリソースの削除（VMとTPUの両方が削除されます）

ctpu delete --name=tpu-v3-32-test

削除には1〜2分程度かかります。
TPUのみを削除したい場合は、コンソールからTPUのみを削除することもできます。

GPUでのトレーニング

比較のためにGPUを使ったResNetの学習も行いました。GPUはGCP上でも提供されており、誰でも利用することができます。（参考)

我々がTPUを利用したときに参照したGitHub内にあったこちらの資料を参考にして、GPU V100を8台使い、総計算量がTPUと同じになるようにパラメータをチューニングし、同様にトレーニングを行いました。

実験結果（トレーニング時間と費用）

各アーキテクチャで今回のモデル・データ量についてトレーニングした結果は以下です。

*1 1$ = 107.68円

GPUを実際に買うとなると高いですがGCPの従量課金サービスであるCompute EngineのGPUを使えば、使った分だけしか費用がかかりません。今回の検証はGCP上で8台のGPU NVIDIA Tesla V100を利用しました。
結果としてはTPU Pod v3-32の1.5倍のトレーニング時間、1.7倍以上のトレーニング費用がかかりました。トレーニング時間、費用ともにTPUが圧勝しています。
以上から、Cloud TPUは非常に費用対効果に優れていると言えるでしょう。

TPU Pod の性能とコストまとめ

Googleが開発したTPUに関しての説明と、GCPプロダクトであるCloud TPUのv3 Podとv2 PodでResNetをトレーニングして性能とコストの比較を行いました。
TPUのセットアップは非常に容易で、ハードウェアをいじる必要がなくコマンドのみで完結します。また性能も自由に選ぶことができます。パフォーマンスが足りなければPodを使うと良いでしょう。

ちなみに、TPU Podは2019年8月時点では未だβですので挙動がおかしい部分がある可能性があります。本検証においても何度か躓くことがありました。怪しい場合はすぐにサポートに連絡して、少なくともTPUはシャットダウンする等の対応をした方が良いでしょう。

とはいえ、これほど強力なプロセッサであるTPUをこれほど安く利用できるのは驚きでした。モデルのトレーニング時間がかかりすぎて困っている＆予算が限られている方はぜひCloud TPUの利用を考えてみてはいかがでしょうか？

おまけ

本実験ではトレーニング自体は各アーキテクチャでの合計で1万円ほどで実施可能でしたが、実際の検証の過程では9万円費やしました。立ち上げシャットダウンや準備にかかる時間、ベースとした試行用のソースコードやベータ版であるがゆえのバグなどを踏み抜いたことが原因です。
VM内にあるResNetのモデルのソースコードはGitHubにあり、常に修正と改良がされています。検証期間中にもソースコードが変わってドキュメントと引数が異なっていたり（修正済み）とかなり苦労しました。そのため、起動中のお金が発生する中で、ベータ版ゆえのドキュメントの少なさに試行錯誤して色々と検証していたからです。特にTensorFlow ver 1.13以下で起きるトレーニング途中でハングしてしまうバグには非常に苦労させられました。ハング中もTPUは起動しているので、その時間分のお金は請求されます。失敗している理由が不明なままお金が発生するという状況に精神力がゴリゴリ削られました。。。

ただ、（公式ドキュメントではCloud TPUはTensorFlow 1.13まで対応していないことになっていますが）最近デフォルトでTensorFlow 1.14のTPU Podが立つようになったので、もうこのバグに悩まされる心配はないです。僕らのサポートへの直訴が効いた結果でしょう。