将 TPU 与 PyTorch 一起使用

Posted 2023-03-12

【中文标题】将 TPU 与 PyTorch 一起使用

【英文标题】：Using TPUs with PyTorch

【发布时间】：2019-03-10 04:04:59

【问题描述】：

我正在尝试使用 Colab 的 Google Cloud 的 TPU。我能够按照教程使用 Tensorflow 做到这一点。

有人知道是否可以使用 PyTorch 来使用 TPU？ 如果是这样，我该怎么做？你有什么例子吗？

【参考方案1】：

目前，无法将 Cloud TPU 与 PyTorch 结合使用，因为它是专为 Tensorflow 设计的。

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

【讨论】：

我明白了！非常感谢！

【参考方案2】：

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

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

【讨论】：

非常感谢您的辛勤工作！这是一个重要的里程碑！

【参考方案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)

【讨论】：

我尝试运行文档协作，这是我得到的错误：module 'torch_xla.core.xla_model' has no attribute 'rendezvous'