在PyTorch中使用TPU

Question

我正在尝试使用Colab的Google Cloud TPU。通过使用Tensorflow，我可以按照本教程进行操作。

有人知道是否可以通过PyTorch使用TPU？ 如果可以，我该怎么办？你有什么例子吗？

Answer 1

由于PyTorch专为Tensorflow设计，因此目前无法将Cloud TPU与PyTorch一起使用。

但是，据this product news posted three days ago in the Google Cloud blog称，“ Google TPU团队的工程师正在与PyTorch核心开发人员积极合作，以将PyTorch连接到云TPU。”

Answer 2

到目前为止，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)

Answer 3

查看我们的存储库pytorch/xla，您可以在其中开始在TPU上训练PyTorch模型。

此外，您甚至可以通过这些Colab notebooks在带有PyTorch的Colab上使用免费的TPU。