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で実行可能

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

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を学習させてみる

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

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

15.
モデルの並列化
• Train
• Valid

16.
Train loop

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

18.
pytorch-lightningに
書き換えてみた

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

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

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

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でログを出力できる
• 簡潔に書ける一方でカスタマイズ性も高い