如何在 PyTorch 中使用 TPU？答案

【问题标题】：How to use TPUs with PyTorch?如何在 PyTorch 中使用 TPU？
【发布时间】：2020-05-17 05:36:06
【问题描述】：

我正在尝试使用 pytorch_xla 使用 TPU，但它在 _XLAC 中显示导入错误。

!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py
!python pytorch-xla-env-setup.py --version $VERSION

import torch_xla
import torch_xla.core.xla_model as xm

ImportError                               Traceback (most recent call last)
<ipython-input-60-6a19e980152f> in <module>()
----> 1 import torch_xla
      2 import torch_xla.core.xla_model as xm

/usr/local/lib/python3.6/dist-packages/torch_xla/__init__.py in <module>()
     39 import torch
     40 from .version import __version__
---> 41 import _XLAC
     42 
     43 _XLAC._initialize_aten_bindings()

ImportError: /usr/local/lib/python3.6/dist-packages/_XLAC.cpython-36m-x86_64-linux-gnu.so: undefined symbol: _ZN2at6native6einsumENSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEEN3c108ArrayRefINS_6TensorEEE

【问题讨论】：

标签： python pytorch tpu

【解决方案1】：

确保您使用的是正确版本的 pytorch-xla 和 Python（3.6.9 运行良好）：

curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py
python pytorch-xla-env-setup.py --version 20200325

检查您是否指定了如何访问 TPU。您可能需要设置“XRT_TPU_CONFIG”或“COLAB_TPU_ADDR”，具体取决于您的环境。

类似：

export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"

或者：

export COLAB_TPU_ADDR="10.16.26.36:8676"

这里是详细描述：https://github.com/pytorch/xla/blob/master/README.md 和示例https://cloud.google.com/tpu/docs/tutorials/transformer-pytorch

另外，这里是 PyTorch 团队创建的 Google Colab 笔记本，我刚刚测试过，它可以正常工作，无需任何更改： https://colab.research.google.com/github/pytorch/xla/blob/master/contrib/colab/getting-started.ipynb

此笔记本将向您展示如何：

在 Colab 上安装 PyTorch/XLA，让您可以将 PyTorch 与 TPU 结合使用。

在 TPU 上运行基本的 PyTorch 函数。

在 TPU 上运行 PyTorch 模块和 autograd。

在 TPU 上运行 PyTorch 网络。

您可能希望遵循其中一个示例并尝试重现该问题。祝你好运！

【讨论】：

【解决方案2】：

请试试这个：

!pip uninstall -y torch
!pip install torch==1.8.2+cpu -f https://download.pytorch.org/whl/lts/1.8/torch_lts.html
!pip install -q cloud-tpu-client==0.10 https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-1.8-cp37-cp37m-linux_x86_64.whl
import torch_xla

它对我有用。

来源：googlecolab/colabtools#2237

【讨论】：

正确！确保 PyTorch 和 PyTorch/XLA 之间的包版本是一致的，并且与 Cloud TPU Colabs 上的建议一致。