[DLHacks]pytorch - segmentation を TPU で実行してみた / pytorch - lightning で書き換えてみた

pytorch-segmentationを TPUで実行してみた/ pytorch-lightningで書き換えてみた 東京大学大学院 情報理工学系研究科 電子情報学専攻 坂井・入江研 D1 谷合 廣紀

今日の内容 • セマンティックセグメンテーションを題材に • PyTorchでTPUを使ってみる • pytorch-lightningで学習コードを書き直す

pytorch-segmentation • https://github.com/nyoki-mtl/pytorch-segmentation • PyTorchで書いたセマンティックセグメンテーション用のコード • Deeplab V3+, MobileNet V2, Unetなどが使える

• https://github.com/nyoki-mtl/pytorch-segmentation

TPUで実行してみた

TPUとは • Googleが開発したディープラーニング用のプロセッサ • Google ColaboratoryやGoogle Cloud Platformで使える • フレームワークとしえはKeras/Tensorflowが対応

PyTorchからもTPUを使ってみたい • 去年10月くらいにPyTorchも今後TPUに対応していくことが発表された • • https://cloud.google.com/blog /products/ai-machinelearning /introducing-pytorch-across-google-cloud https://github.com/pytorch/xla で開発が進められている • XLA(Accelerated Linear Algebra)はTensorflow用のコンパイラ • XLAでコンパイルした中間言語はTPUで実行可能

• https://cloud.google.com/blog/products/ai-machine-learning/introducing-pytorch-across-google-cloud
• https://github.com/pytorch/xla

pytorch/xla • githubのページに行くとチュートリアルがある • 実行してみよう！

環境 • Google Cloud Platformで実験 • インスタンス • TPU v3x8 (第三世代のTPUが8台) • n1-standard-16 (vCPU x16, Memory 60GB)

TPUの起動 (コマンド例) $ gcloud compute tpus create tpu-node1 \ --zone=us-central1-a \ --network=default \ --range=10.0.101.0 \ Pytorch用のソフトウェア --version=pytorch-0.1 \ --accelerator-type=v3-8

CPUの起動 (コマンド例) $ gcloud compute instances create master-vm \ --zone=us-central1-a \ --machine-type=n1-standard-16 \ --image=debian-9-torch-xla-v20190817 \ --image-project=ml-images \ --boot-disk-size=100GB \ Pytorch用のイメージ --boot-disk-type=pd-ssd

MNIST実行 $ gcloud compute ssh master-vm (vm)$ export XRT_TPU_CONFIG="tpu_worker;0;10.0.101.2:8470“ (vm)$ conda activate pytorch-0.1 (pytorch-0.1)$ cd /usr/share/torch-xla-0.1/pytorch/xla (pytorch-0.1)$ python test/test_train_mnist.py

pytorch-segmentationを書き換える • もとのコード https://github.com/nyoki-mtl/pytorch-segmentation/blob/master/src /train.py • TPU用のコード https://github.com/nyoki-mtl/pytorch-segmentation/blob/tpu/src /train_tpu.py • UNetでcityscapesを学習させてみる

• https://github.com/nyoki-mtl/pytorch-segmentation/blob/master/src/train.py
• https://github.com/nyoki-mtl/pytorch-segmentation/blob/tpu/src/train_tpu.py

ライブラリ • TPU用の環境変数 • torch-xlaのインポート

TPUの設定 • 使用するデバイスの設定 • max_devices: TPUをいくつ使うか指定でき、Noneなら全部使う

モデルの並列化 • Train • Valid

Train loop

所感 • 思ったより簡単にTPU実行可能なコードに書き換えることができた • まだできていない部分も多い • • Dilated convが実行できなかったので、deeplab v3+は動かせなかった • model.eval()をするとBatchNormの挙動がおかしくなっているように感じた 学習に時間がかかる • NVIDIA TITAN Vで3分の学習がTPU v3x8で20分ほどかかった • TPUの強みを活かせていない or 別のところにボトルネックがあるかも

pytorch-lightningに 書き換えてみた

pytorch-lightningとは • PyTorchのラッパーで効率的に実験コードを書くことができるライブラリ • Tensorflowにおけるkerasのような立ち位置 • 最近、Pytorch公式のエコシステムのひとつに

pytorch-lightningの書き方 • pl.LightningModuleを継承したモデルを書いて

pytorch-lightningの書き方 • Trainerで学習を回すだけ

pytorch-segmentationを書き換える • もとのコード • • • https://github.com/nyoki-mtl/pytorch-segmentation/blob/master/src/train.py Pytorch-lightningのコード • https://github.com/nyoki-mtl/pytorch-segmentation/blob/pytorch-lightning/src/pl_model.py • https://github.com/nyoki-mtl/pytorch-segmentation/blob/pytorch-lightning/src/pl_train.py Deeplabv3+でcityscapesを学習させてみる

• https://github.com/nyoki-mtl/pytorch-segmentation/blob/master/src/train.py
• https://github.com/nyoki-mtl/pytorch-segmentation/blob/pytorch-lightning/src/pl_model.py
• https://github.com/nyoki-mtl/pytorch-segmentation/blob/pytorch-lightning/src/pl_train.py

所感 • コードがきれいにまとまって可読性が向上した • 機能が多い • • 簡単にマルチGPUや混合精度演算を実行できる • Tensorboardでログを出力できる 簡潔に書ける一方でカスタマイズ性も高い