Pytorchで始めるはじめてのTPU

社内の輪講での発表資料です。
TPUについての簡単な説明とPytorchでTPUを用いてImageNetを学習させたときの肌感をまとめました。
TPUについて初心者なので、間違ったことがあればご指摘いただきたいですm(_ _)m

Transcript

1. Pytorchで始めるはじめてのTPU AIシステム部 データサイエンス第⼀グループ Naoto Shimakoshi 2019/01/ 23
2. ⾃⼰紹介 2 • 名前 • 島越 直⼈ (シマコシ ナオト) • よくトリゴエと間違えられますがシマコシです。 • Twitter • @nt_4o54 • 経歴 • 奈良県出⾝ • KUの機械系⼤学院 卒 • 2019/04 ~ DeNA新卒⼊社 • Kaggle • @shimacos (⾦ 1, ソロ銀 2) • Kaggle Master • ⾊々なドメインのデータに触れるのが好き
3. ⽬的 • Kaggleや業務で⼤きなデータセットを⽤いて訓練させる際に実験サイクルを速くしたい。 • コンペ終盤の焦ってる時など、我慢できずに⼿動枝刈りしてしまいがち。 • ImageNetのpretrainをよく使うけれど、BatchNormなどのmoduleは固定されている。 • 実はGroupNormとかでpretrainした⽅が良かったりして、そのweightを所持していると 有利だったりする︖︖ • (BatchNormだけをreplaceすることも考えられるが、⼀旦考えない) • (ImageNetのpretrainは学習を⾼速にしているだけという話もあるが、⼀旦考えない) • 最新のまだweightが公開されていないmodelとか試したい。 3 とにかく⼤きなモデルを⾼速に実験したい！
4. ⽬的 • Kaggleや業務で⼤きなデータセットを⽤いて訓練させる際に実験サイクルを速くしたい。 • コンペ終盤の焦ってる時など、我慢できずに⼿動枝刈りしてしまいがち。 • ImageNetのpretrainをよく使うけれど、BatchNormなどのmoduleは固定されている。 • 実はGroupNormとかでpretrainした⽅が良かったりして、そのweightを所持していると 有利だったりする︖︖ • (BatchNormだけをreplaceすることも考えられるが、⼀旦考えない) • (ImageNetのpretrainは学習を⾼速にしているだけという話もあるが、⼀旦考えない) • 最新のまだweightが公開されていないmodelとか試したい。 4 とにかく⼤きなモデルを⾼速に実験したい！ → 今回はImageNetをTPUを⽤いてtrainさせた肌感を伝えようと思います
5. TPU (Tensor Computing Unit) とは • Googleが開発したプロセッサ • ⼤規模なパイプライン処理により、⼤量の⾏列演算をわずかなメモリアクセスで実現 5 https://storage.googleapis.com/nexttpu/index.html
6. どういう仕組み︖(Tensorflowでの場合) 6 https://cloud.google.com/tpu/docs/system-architecture プログラムを計算グラフ化し、 Cloud TPU上のサーバに送信 ① TFRecord形式などでGCSからデータを読み込む ② TPU上で実⾏可能な部分とCPU上で実⾏する部分を分割 TPU上で実⾏する部分をXLAコンパイラでバイナリ化する ③
7. TPUのバージョン 7 https://cloud.google.com/tpu/?hl=ja • Cloud TPU, TPU Podの⼆種類とv2とv3の⼆つのバージョンがある • TPU Podは Cloud TPUを専⽤の⾼速ネットワークで相互に接続したTPUデバイスの クラスタ
8. TPUのバージョン 8 TPU Pod https://cloud.google.com/tpu/docs/tpus Preemptible ＄1.35 / hour (=V100 x1 と同じくらい) Preemptible ＄2.40 / hour
9. どういう時に有効か • ⾏列計算が多くを占めるモデル • CPU上で実⾏しなければいけない演算がトレーニングループ内にないモデル • TPU上で実⾏できないTensorflow演算 • バッチサイズを⼤きくとりたい⼤規模なモデル • ラベルノイズが⼤きい状況など︖ • ⾼精度の演算を必要としない (倍精度演算は適していない) • 頻繁な分岐を必要としない 9 https://cloud.google.com/tpu/docs/tpus
10. TensorflowにおけるTPU実⾏コード • 基本は TPUEstimator というものを⽤いる。 • model_fn 内でmodelを定義。TRAIN, EVAL, PREDICT時などで処理を変える必要あり。 • Estimatorに⼊⼒する input_fn を tf.data APIで定義。 • (features, label) の出⼒にする必要あり。TFRecordから読み込む事が多い。 • 慣れていないとめんどくさそう。。。(公式のExample (↓)も読みにくい) 10 https://github.com/tensorflow/tpu/blob/master/models/official/resnet/resnet_main.py
11. Tensorflow1系ではできれば書きたくない • ⽣Tensorflowで開発するのは時間がかかりがち。 • TFRecordは⼀回作ってしまえば便利だが、作るのがめんどくさい。 • 読み込み⽅もめんどくさい。 • Estimatorは⾼レベルAPIを謳っているが（個⼈的には）逆に複雑。 • input_fnとmodel_fnの繋ぎ込みがややこしい。 • Kaggleでやるような複雑な処理を記述しづらい。 11 できることならPytorchでTPU使いたい！ (おそらく多くの⼈の声)
12. PytorchでのTPU実⾏コード① (MultiThreading) • torch_xla.distributed.data_parallel を⽤いる場合 12 DistributedSamplerを⽤いて samplerを定義 (DistriburtedDataParallelの時と同様) TPUのdeviceを取得 [‘xla:1’, ‘xla:2’, …, ‘xla:8’] nn.DataParallelと同様の書き⽅で 各TPUコアにmodelを分散させる deviceはloopが⾛っているTPU core contextは各TPU coreが持つoptimizer やlr_schedulerを格納したもの TPU上にあるmodelをCPUに戻して から保存する ※保存した後はTPUに送り直す
13. PytorchでのTPU実⾏コード① (MultiThreading) • train_loop_fn の例 13 contextは各device毎に定義されており、 各deviceに紐づく属性として optimizerなどを格納している。 getattr_orはos.environ.getの様なもの (http://pytorch.org/xla/_modules/torch_ xla/distributed/data_parallel.html) xm.optimizer_stepでラップすることで、異 なるdevice間の 勾配をまとめて平均する
14. PytorchでのTPU実⾏コード② (MultiProcessing) • torch_xla.distributed.parallel_loader と torch_xla.distributed.xla_multiprocessing を⽤いる場合 14 _mp_fn を num_cores分だけ MultiProcessingに並列処理を⾏う ↓ TPU device分だけprocessを作成 indexは processのindex
15. PytorchでのTPU実⾏コード② (MultiProcessing) • train_imagenet の中⾝ (各device内で⾛る処理を記述すれば良い) 15 Processが存在するdeviceを取得し、 modelをdevice上に移動する Optimizerなどは各deviceで 同⼀のものを使うのでloop_fnの外側に定義 pl.ParallelLoaderでloaderを並列に処理 (DataLoaderのnum_workerと 何が違うかよくわかっていないが速い) per_device_loader でprocessに対応する loaderのみを取得 DataParallel と同じ
16. TFに⽐べてめっちゃ簡単に実装できそう︕︖︕︖ 16
17. 落とし⽳ • PytorchのdockerのversionとTPU側のバージョンを合わせないとめちゃくちゃ遅くなる • Pytorchは最新のstableバージョンがr0.5とタグづけされているが、Pytorchのversionが 0.5な訳ではない。 • 最新のpytorchを使いたかったので、nightly_YYYYMMDDバージョンを使ったが遅す ぎて死んだ。 • suffix無しのnightlyバージョンだといけそう • ちゃんと調べたら、r0.5のdockerに⼊っているPytorchは1.3だった。 • TFのTPUサーバー上でデータを読み込みキューするといった機能がPytorchにはない。 • VM上のリソースを⽤いてDataLoaderを処理するため、VMも⾼価になりがち。 • 今回の実験で使ったのはN1-highmem-94 • XLA_USE_BF16=1 をしないとTPUにおける混合精度にならず遅い。 17
18. 基本的な使い⽅ 18 • 提供されているDockerimageを使った⽅が良い。 • 最新のstable版は gcr.io/tpu-pytorch/xla:r0.5 • nightlyバージョンを使って問題が起きたときはr0.5を⽤いる運⽤が良さそう • TPUとVMを⽴ち上げる • ctpu upを⽤いるかUI上で⽴ち上げる。 ※ TPUのソフトウェアバージョンに注意 • VMとTPUのリージョンは同じにする。 • 内部IPを⽤いて $ export TPU_IP_ADDRESS= • ImageNetをVMにダウンロードしておく。 • pytorch/xla を⼿元にcloneしておく。(dockerコンテナ内の.py書き換え⽤)
19. • とりあえずはこれで動く。Modelだけを書き換えれば汎⽤性⾼そう。 基本コマンド (MultiProcessingの例) 19 Code更新⽤ データをマウント Log吐き場 GCSなどのAuth⽤
20. モニタリング⽅法 • cloud-tpu-profiler をinstallする • pip install cloud-tpu-profiler (ただし、これはtensorflow2系だと動かない) • capture_tpu_profile --tpu=pytorch-trainer-r5 --monitoring_level=1 20 ここを⾒ながらbatch_sizeなど チューニングするといいかも
21. TensorflowとPytorchの速度⽐較 • N1-highmem-94をVMに⽤いた。 • Resnet50で精度76%に到達するまでの時間を計測 (batch_size=1024) 21 TPU version resnet50 Pytorch (MultiThreading) TPU v3-8 17 hour Pytorch (MultiProcessing) TPU v3-8 10 hour TF1.xx TPU v3-8 9.0 hour
22. 感想 • TFは速いが開発が⾟い。 • Pytorchの⽅が圧倒的に実装が楽ではあるが遅い。 • 動的に(?)グラフをCompileするため最初の⽅は遅く、2epoch⽬くらいから速くなる。 • Cloud TPU上で直接GCSからデータを読み込めない。 • tf.data APIの様に画像読み込みやAugmentationをグラフとして構築できない。 • GPUインスタンスの様にPreemptibleでも全然落ちないのでそれだけでも嬉しい。 • Tensorboardにtrain_loop_fnからlogを送ろうとすると動かなくなる。闇。 • たまにDataParallelだとMultiThreadingがhangして動かなくなる。闇。 • https://github.com/pytorch/xla/issues/821 • MultiProcessingだと何故かGCSにlogを送れない。闇。 • nightlyバージョンだと送れた 22
23. 感想 • 現時点ではなんだかんだTensorflow EstimatorがTPU対応がリッチなので強そう。 • この資料を書いている内にEstimatorでも良くね︖という気持ちになってきた。 • TF2.1からはtf.kerasもexperimental supportになったらしい。 • https://github.com/tensorflow/tensorflow/releases • PytorchやTF2.0のCloud TPU関連の資料は少なく、コード例もほぼ公式のものしかない。 • TPUを使って複雑なことをしたいときは強い意志が必要。 • とりあえずImageNetの例におけるmodelとdatasetを⾃作のものに変換すれば使えそ う︖ • 今回の実験に⽤いたコードは以下のレポジトリにあるので興味のある⽅はどうぞ。 • https://github.com/shimacos37/tpu_experiment 23
24. まとめ • Tensorflow • メリット • 公式のサポートがリッチ。 • Pytorchより⾼速。 • デメリット • 慣れていないと開発に時間がかかる。 • Pytorch • MultiProcessingを使うのが肌感的には安定している。 • メリット • 開発がやりやすい。既存のコード資産も使える。 • デメリット • TFに⽐べると遅い。 • VMも⾼機能を要求されるのでその分費⽤がかかる。 24