Pytorchで始めるはじめてのTPU

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

Transcript

1. Pytorchで始めるはじめてのTPU
   AIシステム部 データサイエンス第⼀グループ
   Naoto Shimakoshi
   2019/01/ 23
   

   View Slide

2. ⾃⼰紹介
   2
   • 名前
   • 島越 直⼈ (シマコシ ナオト)
   • よくトリゴエと間違えられますがシマコシです。
   • Twitter
   • @nt_4o54
   • 経歴
   • 奈良県出⾝
   • KUの機械系⼤学院 卒
   • 2019/04 ~ DeNA新卒⼊社
   • Kaggle
   • @shimacos (⾦ 1, ソロ銀 2)
   • Kaggle Master
   • ⾊々なドメインのデータに触れるのが好き
   

   View Slide

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

   View Slide

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

   View Slide

5. TPU (Tensor Computing Unit) とは
   • Googleが開発したプロセッサ
   • ⼤規模なパイプライン処理により、⼤量の⾏列演算をわずかなメモリアクセスで実現
   5
   https://storage.googleapis.com/nexttpu/index.html
   

   View Slide

6. どういう仕組み︖(Tensorflowでの場合)
   6
   https://cloud.google.com/tpu/docs/system-architecture
   プログラムを計算グラフ化し、
   Cloud TPU上のサーバに送信
   ①
   TFRecord形式などでGCSからデータを読み込む
   ②
   TPU上で実⾏可能な部分とCPU上で実⾏する部分を分割
   TPU上で実⾏する部分をXLAコンパイラでバイナリ化する
   ③
   

   View Slide

7. TPUのバージョン
   7
   https://cloud.google.com/tpu/?hl=ja
   • Cloud TPU, TPU Podの⼆種類とv2とv3の⼆つのバージョンがある
   • TPU Podは Cloud TPUを専⽤の⾼速ネットワークで相互に接続したTPUデバイスの
   クラスタ
   

   View Slide

8. TPUのバージョン
   8
   TPU Pod https://cloud.google.com/tpu/docs/tpus
   Preemptible
   ＄1.35 / hour
   (=V100 x1
   と同じくらい)
   Preemptible
   ＄2.40 / hour
   

   View Slide

9. どういう時に有効か
   • ⾏列計算が多くを占めるモデル
   • CPU上で実⾏しなければいけない演算がトレーニングループ内にないモデル
   • TPU上で実⾏できないTensorflow演算
   • バッチサイズを⼤きくとりたい⼤規模なモデル
   • ラベルノイズが⼤きい状況など︖
   • ⾼精度の演算を必要としない (倍精度演算は適していない)
   • 頻繁な分岐を必要としない
   9
   https://cloud.google.com/tpu/docs/tpus
   

   View Slide

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
    

    View Slide

11. Tensorflow1系ではできれば書きたくない
    • ⽣Tensorflowで開発するのは時間がかかりがち。
    • TFRecordは⼀回作ってしまえば便利だが、作るのがめんどくさい。
    • 読み込み⽅もめんどくさい。
    • Estimatorは⾼レベルAPIを謳っているが（個⼈的には）逆に複雑。
    • input_fnとmodel_fnの繋ぎ込みがややこしい。
    • Kaggleでやるような複雑な処理を記述しづらい。
    11
    できることならPytorchでTPU使いたい！
    (おそらく多くの⼈の声)
    

    View Slide

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に送り直す
    

    View Slide

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間の
    勾配をまとめて平均する
    

    View Slide

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
    

    View Slide

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 と同じ
    

    View Slide

16. TFに⽐べてめっちゃ簡単に実装できそう︕︖︕︖
    16
    

    View Slide

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
    

    View Slide

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書き換え⽤)
    

    View Slide

19. • とりあえずはこれで動く。Modelだけを書き換えれば汎⽤性⾼そう。
    基本コマンド (MultiProcessingの例)
    19
    Code更新⽤
    データをマウント
    Log吐き場
    GCSなどのAuth⽤
    

    View Slide

20. モニタリング⽅法
    • cloud-tpu-profiler をinstallする
    • pip install cloud-tpu-profiler (ただし、これはtensorflow2系だと動かない)
    • capture_tpu_profile --tpu=pytorch-trainer-r5 --monitoring_level=1
    20
    ここを⾒ながらbatch_sizeなど
    チューニングするといいかも
    

    View Slide

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
    

    View Slide

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
    

    View Slide

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
    

    View Slide

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

    View Slide