Как использовать TPU с PyTorch?

Я пытаюсь использовать TPU с помощью pytorch_xla, но он показывает ошибку импорта в _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
person sharma.37    schedule 17.05.2020    source источник

Ответы (1)


arrow_upward
1
arrow_downward

  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
  1. Убедитесь, что вы указали, как получить доступ к TPU. Вы можете настроить «XRT_TPU_CONFIG» или «COLAB_TPU_ADDR» в зависимости от вашей среды.

Что-то вроде:

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

Or: 

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.

Кроме того, вот блокнот Google Colab, созданный командой PyTorch, я только что протестировал его, он работает нормально без каких-либо изменений: https://colab.research.google.com/github/pytorch/xla/blob/master/contrib/colab/getting-started.ipynb

Этот блокнот покажет вам, как:

  • Установите PyTorch / XLA на Colab, что позволяет использовать PyTorch с TPU.
  • Запускайте базовые функции PyTorch на TPU.
  • Запускайте модули PyTorch и выполняйте автоматическую настройку на TPU.
  • Запускайте сети PyTorch на TPU.

Вы можете последовать одному из примеров и попытаться воспроизвести проблему. Удачи!

person AI Mechanic    schedule 25.05.2020