[DLHacks]pytorch - segmentation を TPU で実行してみた / pytorch - lightning で書き換えてみた
Download as PPTX, PDF2 likes1,784 views
[DLHacks]pytorch - segmentation を TPU で実行してみた / pytorch - lightning で書き換えてみた
- 1. pytorch-segmentationを TPUで実行してみた/ pytorch-lightningで書き換えてみた 東京大学大学院 情報理工学系研究科 電子情報学専攻 坂井・入江研 D1 谷合 廣紀
- 2. 今日の内容 • セマンティックセグメンテーションを題材に • PyTorchでTPUを使ってみる • pytorch-lightningで学習コードを書き直す
- 3. pytorch-segmentation • https://github.com/nyoki-mtl/pytorch-segmentation • PyTorchで書いたセマンティックセグメンテーション用のコード • Deeplab V3+, MobileNet V2, Unetなどが使える
- 4. TPUで実行してみた
- 5. TPUとは • Googleが開発したディープラーニング用のプロセッサ • Google ColaboratoryやGoogle Cloud Platformで使える • フレームワークとしえはKeras/Tensorflowが対応
- 6. PyTorchからもTPUを使ってみたい • 去年10月くらいにPyTorchも今後TPUに対応していくことが発表された • https://cloud.google.com/blog/products/ai-machine- learning/introducing-pytorch-across-google-cloud • https://github.com/pytorch/xla で開発が進められている • XLA(Accelerated Linear Algebra)はTensorflow用のコンパイラ • XLAでコンパイルした中間言語はTPUで実行可能
- 8. 環境 • Google Cloud Platformで実験 • インスタンス • TPU v3x8 (第三世代のTPUが8台) • n1-standard-16 (vCPU x16, Memory 60GB)
- 9. TPUの起動 (コマンド例) $ gcloud compute tpus create tpu- node1 --zone=us-central1-a --network=default --range=10.0.101.0 --version=pytorch-0.1 --accelerator-type=v3-8 Pytorch用のソフトウェア
- 10. 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 --boot-disk-type=pd-ssd Pytorch用のイメージ
- 11. 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
- 12. 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を学習させてみる
- 14. TPUの設定 • 使用するデバイスの設定 • max_devices: TPUをいくつ使うか指定でき、Noneなら全部使う
- 16. Train loop
- 17. 所感 • 思ったより簡単にTPU実行可能なコードに書き換えることができた • まだできていない部分も多い • Dilated convが実行できなかったので、deeplab v3+は動かせなかった • model.eval()をするとBatchNormの挙動がおかしくなっているように感じた • 学習に時間がかかる • NVIDIA TITAN Vで3分の学習がTPU v3x8で20分ほどかかった • TPUの強みを活かせていない or 別のところにボトルネックがあるかも
- 19. pytorch-lightningとは • PyTorchのラッパーで効率的に実験コードを書くことができるライブラリ • Tensorflowにおけるkerasのような立ち位置 • 最近、Pytorch公式のエコシステムのひとつに
- 22. 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を学習させてみる
- 23. 所感 • コードがきれいにまとまって可読性が向上した • 機能が多い • 簡単にマルチGPUや混合精度演算を実行できる • Tensorboardでログを出力できる • 簡潔に書ける一方でカスタマイズ性も高い