如何在Google Colab为PyTorch和PyTorch Lightning配置TPU？安装报错求助

问题描述 #

我按照PyTorch Lightning官方指南配置TPU，执行了以下命令：

! pip install --quiet "pytorch-lightning>=1.4" "setuptools==59.5.0" "torch>=1.8" "torchmetrics>=0.7" "torchvision" "ipython[notebook]"

! pip install cloud-tpu-client==0.10 https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-1.8-cp37-cp37m-linux_x86_64.whl

但遇到错误：

查看索引：https://pypi.org/simple、https://us-python.pkg.dev/colab-wheels/public/simple/
ERROR: torch_xla-1.8-cp37-cp37m-linux_x86_64.whl is not a supported wheel on this platform.

解决方案 #

• 匹配Python版本与wheel包：报错的wheel包是为Python 3.7（cp37）编译的，若当前环境使用更高版本Python（如3.8/3.9）则会不兼容。先执行python --version确认Python版本，再替换为对应版本的torch_xla wheel包链接。

• 使用自动适配的安装脚本：避免手动指定固定wheel链接，改用官方脚本自动匹配平台和Python版本：

  ! pip install cloud-tpu-client==0.10
  ! 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 1.8 --apt-packages libomp5 libopenblas-dev
  

• 确认运行环境配置：如果在Colab中操作，需先切换到TPU运行时（菜单栏「代码执行程序」→「更改运行时类型」，选择TPU）。Colab的TPU环境会自动适配依赖，也可直接使用平台提供的初始化代码完成配置。

内容的提问来源于stack exchange，提问作者Grateful_programmer