Lock in $30 Savings on PRO—Offer Ends Soon! ⏳

Pytorchで始めるはじめてのTPU

shimacos
January 23, 2020

 Pytorchで始めるはじめてのTPU

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

shimacos

January 23, 2020
Tweet

More Decks by shimacos

Other Decks in Programming

Transcript

  1. ⾃⼰紹介 2 • 名前 • 島越 直⼈ (シマコシ ナオト) •

    よくトリゴエと間違えられますがシマコシです。 • Twitter • @nt_4o54 • 経歴 • 奈良県出⾝ • KUの機械系⼤学院 卒 • 2019/04 ~ DeNA新卒⼊社 • Kaggle • @shimacos (⾦ 1, ソロ銀 2) • Kaggle Master • ⾊々なドメインのデータに触れるのが好き
  2. ⽬的 • Kaggleや業務で⼤きなデータセットを⽤いて訓練させる際に実験サイクルを速くしたい。 • コンペ終盤の焦ってる時など、我慢できずに⼿動枝刈りしてしまいがち。 • ImageNetのpretrainをよく使うけれど、BatchNormなどのmoduleは固定されている。 • 実はGroupNormとかでpretrainした⽅が良かったりして、そのweightを所持していると 有利だったりする︖︖

    • (BatchNormだけをreplaceすることも考えられるが、⼀旦考えない) • (ImageNetのpretrainは学習を⾼速にしているだけという話もあるが、⼀旦考えない) • 最新のまだweightが公開されていないmodelとか試したい。 3 とにかく⼤きなモデルを⾼速に実験したい!
  3. ⽬的 • Kaggleや業務で⼤きなデータセットを⽤いて訓練させる際に実験サイクルを速くしたい。 • コンペ終盤の焦ってる時など、我慢できずに⼿動枝刈りしてしまいがち。 • ImageNetのpretrainをよく使うけれど、BatchNormなどのmoduleは固定されている。 • 実はGroupNormとかでpretrainした⽅が良かったりして、そのweightを所持していると 有利だったりする︖︖

    • (BatchNormだけをreplaceすることも考えられるが、⼀旦考えない) • (ImageNetのpretrainは学習を⾼速にしているだけという話もあるが、⼀旦考えない) • 最新のまだweightが公開されていないmodelとか試したい。 4 とにかく⼤きなモデルを⾼速に実験したい! → 今回はImageNetをTPUを⽤いてtrainさせた肌感を伝えようと思います
  4. TPUのバージョン 7 https://cloud.google.com/tpu/?hl=ja • Cloud TPU, TPU Podの⼆種類とv2とv3の⼆つのバージョンがある • TPU

    Podは Cloud TPUを専⽤の⾼速ネットワークで相互に接続したTPUデバイスの クラスタ
  5. どういう時に有効か • ⾏列計算が多くを占めるモデル • CPU上で実⾏しなければいけない演算がトレーニングループ内にないモデル • TPU上で実⾏できないTensorflow演算 • バッチサイズを⼤きくとりたい⼤規模なモデル •

    ラベルノイズが⼤きい状況など︖ • ⾼精度の演算を必要としない (倍精度演算は適していない) • 頻繁な分岐を必要としない 9 https://cloud.google.com/tpu/docs/tpus
  6. 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
  7. 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に送り直す
  8. 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間の 勾配をまとめて平均する
  9. 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
  10. 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 と同じ
  11. 落とし⽳ • 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
  12. 基本的な使い⽅ 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書き換え⽤)
  13. モニタリング⽅法 • cloud-tpu-profiler をinstallする • pip install cloud-tpu-profiler (ただし、これはtensorflow2系だと動かない) •

    capture_tpu_profile --tpu=pytorch-trainer-r5 --monitoring_level=1 20 ここを⾒ながらbatch_sizeなど チューニングするといいかも
  14. 感想 • 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
  15. 感想 • 現時点ではなんだかんだ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
  16. まとめ • Tensorflow • メリット • 公式のサポートがリッチ。 • Pytorchより⾼速。 •

    デメリット • 慣れていないと開発に時間がかかる。 • Pytorch • MultiProcessingを使うのが肌感的には安定している。 • メリット • 開発がやりやすい。既存のコード資産も使える。 • デメリット • TFに⽐べると遅い。 • VMも⾼機能を要求されるのでその分費⽤がかかる。 24