社内の輪講での発表資料です。 TPUについての簡単な説明とPytorchでTPUを用いてImageNetを学習させたときの肌感をまとめました。 TPUについて初心者なので、間違ったことがあればご指摘いただきたいですm(_ _)m
Pytorchで始めるはじめてのTPUAIシステム部 データサイエンス第⼀グループNaoto Shimakoshi2019/01/ 23
View Slide
⾃⼰紹介2• 名前• 島越 直⼈ (シマコシ ナオト)• よくトリゴエと間違えられますがシマコシです。• Twitter• @nt_4o54• 経歴• 奈良県出⾝• KUの機械系⼤学院 卒• 2019/04 ~ DeNA新卒⼊社• Kaggle• @shimacos (⾦ 1, ソロ銀 2)• Kaggle Master• ⾊々なドメインのデータに触れるのが好き
⽬的• Kaggleや業務で⼤きなデータセットを⽤いて訓練させる際に実験サイクルを速くしたい。• コンペ終盤の焦ってる時など、我慢できずに⼿動枝刈りしてしまいがち。• ImageNetのpretrainをよく使うけれど、BatchNormなどのmoduleは固定されている。• 実はGroupNormとかでpretrainした⽅が良かったりして、そのweightを所持していると有利だったりする︖︖• (BatchNormだけをreplaceすることも考えられるが、⼀旦考えない)• (ImageNetのpretrainは学習を⾼速にしているだけという話もあるが、⼀旦考えない)• 最新のまだweightが公開されていないmodelとか試したい。3とにかく⼤きなモデルを⾼速に実験したい!
⽬的• Kaggleや業務で⼤きなデータセットを⽤いて訓練させる際に実験サイクルを速くしたい。• コンペ終盤の焦ってる時など、我慢できずに⼿動枝刈りしてしまいがち。• ImageNetのpretrainをよく使うけれど、BatchNormなどのmoduleは固定されている。• 実はGroupNormとかでpretrainした⽅が良かったりして、そのweightを所持していると有利だったりする︖︖• (BatchNormだけをreplaceすることも考えられるが、⼀旦考えない)• (ImageNetのpretrainは学習を⾼速にしているだけという話もあるが、⼀旦考えない)• 最新のまだweightが公開されていないmodelとか試したい。4とにかく⼤きなモデルを⾼速に実験したい!→ 今回はImageNetをTPUを⽤いてtrainさせた肌感を伝えようと思います
TPU (Tensor Computing Unit) とは• Googleが開発したプロセッサ• ⼤規模なパイプライン処理により、⼤量の⾏列演算をわずかなメモリアクセスで実現5https://storage.googleapis.com/nexttpu/index.html
どういう仕組み︖(Tensorflowでの場合)6https://cloud.google.com/tpu/docs/system-architectureプログラムを計算グラフ化し、Cloud TPU上のサーバに送信①TFRecord形式などでGCSからデータを読み込む②TPU上で実⾏可能な部分とCPU上で実⾏する部分を分割TPU上で実⾏する部分をXLAコンパイラでバイナリ化する③
TPUのバージョン7https://cloud.google.com/tpu/?hl=ja• Cloud TPU, TPU Podの⼆種類とv2とv3の⼆つのバージョンがある• TPU Podは Cloud TPUを専⽤の⾼速ネットワークで相互に接続したTPUデバイスのクラスタ
TPUのバージョン8TPU Pod https://cloud.google.com/tpu/docs/tpusPreemptible$1.35 / hour(=V100 x1と同じくらい)Preemptible$2.40 / hour
どういう時に有効か• ⾏列計算が多くを占めるモデル• CPU上で実⾏しなければいけない演算がトレーニングループ内にないモデル• TPU上で実⾏できないTensorflow演算• バッチサイズを⼤きくとりたい⼤規模なモデル• ラベルノイズが⼤きい状況など︖• ⾼精度の演算を必要としない (倍精度演算は適していない)• 頻繁な分岐を必要としない9https://cloud.google.com/tpu/docs/tpus
TensorflowにおけるTPU実⾏コード• 基本は TPUEstimator というものを⽤いる。• model_fn 内でmodelを定義。TRAIN, EVAL, PREDICT時などで処理を変える必要あり。• Estimatorに⼊⼒する input_fn を tf.data APIで定義。• (features, label) の出⼒にする必要あり。TFRecordから読み込む事が多い。• 慣れていないとめんどくさそう。。。(公式のExample (↓)も読みにくい)10https://github.com/tensorflow/tpu/blob/master/models/official/resnet/resnet_main.py
Tensorflow1系ではできれば書きたくない• ⽣Tensorflowで開発するのは時間がかかりがち。• TFRecordは⼀回作ってしまえば便利だが、作るのがめんどくさい。• 読み込み⽅もめんどくさい。• Estimatorは⾼レベルAPIを謳っているが(個⼈的には)逆に複雑。• input_fnとmodel_fnの繋ぎ込みがややこしい。• Kaggleでやるような複雑な処理を記述しづらい。11できることならPytorchでTPU使いたい!(おそらく多くの⼈の声)
PytorchでのTPU実⾏コード① (MultiThreading)• torch_xla.distributed.data_parallel を⽤いる場合12DistributedSamplerを⽤いてsamplerを定義(DistriburtedDataParallelの時と同様)TPUのdeviceを取得[‘xla:1’, ‘xla:2’, …, ‘xla:8’]nn.DataParallelと同様の書き⽅で各TPUコアにmodelを分散させるdeviceはloopが⾛っているTPU corecontextは各TPU coreが持つoptimizerやlr_schedulerを格納したものTPU上にあるmodelをCPUに戻してから保存する※保存した後はTPUに送り直す
PytorchでのTPU実⾏コード① (MultiThreading)• train_loop_fn の例13contextは各device毎に定義されており、各deviceに紐づく属性としてoptimizerなどを格納している。getattr_orはos.environ.getの様なもの(http://pytorch.org/xla/_modules/torch_xla/distributed/data_parallel.html)xm.optimizer_stepでラップすることで、異なるdevice間の勾配をまとめて平均する
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
PytorchでのTPU実⾏コード② (MultiProcessing)• train_imagenet の中⾝ (各device内で⾛る処理を記述すれば良い)15Processが存在するdeviceを取得し、modelをdevice上に移動するOptimizerなどは各deviceで同⼀のものを使うのでloop_fnの外側に定義pl.ParallelLoaderでloaderを並列に処理(DataLoaderのnum_workerと何が違うかよくわかっていないが速い)per_device_loader でprocessに対応するloaderのみを取得DataParallel と同じ
TFに⽐べてめっちゃ簡単に実装できそう︕︖︕︖16
落とし⽳• 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• 提供されている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書き換え⽤)
• とりあえずはこれで動く。Modelだけを書き換えれば汎⽤性⾼そう。基本コマンド (MultiProcessingの例)19Code更新⽤データをマウントLog吐き場GCSなどのAuth⽤
モニタリング⽅法• cloud-tpu-profiler をinstallする• pip install cloud-tpu-profiler (ただし、これはtensorflow2系だと動かない)• capture_tpu_profile --tpu=pytorch-trainer-r5 --monitoring_level=120ここを⾒ながらbatch_sizeなどチューニングするといいかも
TensorflowとPytorchの速度⽐較• N1-highmem-94をVMに⽤いた。• Resnet50で精度76%に到達するまでの時間を計測 (batch_size=1024)21TPU version resnet50Pytorch (MultiThreading) TPU v3-8 17 hourPytorch (MultiProcessing) TPU v3-8 10 hourTF1.xx TPU v3-8 9.0 hour
感想• 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
感想• 現時点ではなんだかんだ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_experiment23
まとめ• Tensorflow• メリット• 公式のサポートがリッチ。• Pytorchより⾼速。• デメリット• 慣れていないと開発に時間がかかる。• Pytorch• MultiProcessingを使うのが肌感的には安定している。• メリット• 開発がやりやすい。既存のコード資産も使える。• デメリット• TFに⽐べると遅い。• VMも⾼機能を要求されるのでその分費⽤がかかる。24