Pytorch+Catalyst+ResNet事前学習済モデルをTPUで動かす

Pytorch

筆者のPytorchの学習状況

KaggleだとPytorchで書かれたコードが多くいい加減やらねば・・・と思い最近UdemyでPytorchの講座を受けていました。

Kerasから入った身としてはKerasよりコード量が多いことがすこしネックでした。

 

モチベーション

Pytorch用のフレームワークを使いながらTPUを使えたらいいのにと思っていたところ、出たので学習済ResNet18をロードしてCIFAR10の学習をしてみました。 

 

実装

以下で公開しています。

colab.research.google.com

github.com

 

今後に期待すること

catalystのバージョンがstabelでないため or Pytochの挙動をいまいちわかっていないためか、colabのランタイムをGPUにした場合、

device = utils.get_device() だと'cpu'しか認識されませんでした。

Pytoch＋TPUの実装はKaggleで公開されていたり、こちらの記事 で紹介されていましたが、torch_xlaのコード改変が必要で少し面倒です。

かったるい事は嫌いなタチなんで、今後はこちらの開発が加速していくことに期待しています。

Pytorch+TPU+(KaggleNotebook or GCP)の構成が普段使いできる展望が見えてきたのが非常にうれしいです。