Upgrade to Pro — share decks privately, control downloads, hide ads and more …

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. 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