PyTorchでのTPUの使用

Question

ColabのGoogle CloudのTPUを使用しようとしています。 Tensorflowを使用して、チュートリアルに従ってそれを行うことができました。

PyTorchを使用してTPUを利用できるかどうか誰かが知っていますか？もしそうなら、どうすればそれを行うことができますか？何か例はありますか？

Answer 1

TPUでPyTorchモデルのトレーニングを開始できるリポジトリ pytorch/xla を確認してください。

また、PyTorchを使用してColabで無料のTPUを使用することもできます Colab notebooks 。

Answer 2

現在、Tensorflow用に特別に設計されているため、PyTorchでCloud TPUを使用することはできません。

しかし、 日前にGoogle Cloudブログに投稿されたこの製品ニュース によると、「GoogleのTPUチームのエンジニアはPyTorchのコアデベロッパーと積極的に協力して、PyTorchをCloud TPUに接続しています」。

Answer 3

現在のところ、 PyTorch Lightning を使用すると、TPUでPyTorchコードを簡単に実行できます（XLAライブラリをインストールする必要があります）。 colab のデモノートブックから：

from pytorch_lightning import Trainer

model = CoolSystem()

# most basic trainer, uses good defaults
trainer = Trainer(num_tpu_cores=8)
trainer.fit(model)