pytorch/xlaで遊んでみた(DevFest Shikoku 2020メモ書き)

DevFest Shikoku 2020のもくもく会でリアルタイムにメモ書きしたものである。

 

類似行為を以前一回やった気もするのだがそのときはPytorch側がNightlyだったかもしれない。

今はPyTorchも1.6となりXLAも正式対応扱いだと思うので安心して使うといい。

 

で,デバイスは何かと言えば貧乏人の無料ツールGoogle ColabのTPUですね。

これを使うためにXLAに手を出すと言ってもいい。

知らない人に簡単に説明すると12時間無料で使えるクラウド上のアクセラレータである。(詳しくはググれ)

 

以前やった時よりドキュメント類が整っている。

ずばり以下を参照すれば(コードが読める人は)ほぼ終わる。

https://colab.research.google.com/github/pytorch/xla/blob/master/contrib/colab/getting-started.ipynb

 

初期設定で重要な三行は以下

import os
assert os.environ['COLAB_TPU_ADDR']

!pip install cloud-tpu-client==0.10 https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-1.6-cp36-cp36m-linux_x86_64.whl

一行目二行目はColabの環境変数の設定。

三行目がtorch_xlaのインストールである。

 

これが終われば以下のリンクのドキュメント通り使えるはずである。

PyTorch on XLA Devices — PyTorch/XLA master documentation

ウチではimport部分で以下のような警告が出たが最終的には通ったようである。

WARNING:root:Waiting for TPU to be start up with version pytorch-1.6...
WARNING:root:Waiting for TPU to be start up with version pytorch-1.6...
WARNING:root:Waiting for TPU to be start up with version pytorch-1.6...
WARNING:root:TPU has started up successfully with version pytorch-1.6

 

その他,学習部が概ね動きそうなんだけど教師データをColab上にアップロードするのが時間的に無理かな。