在Google Colab TPU运行PyTorch Lightning项目出现导入错误如何解决

问题根源 #

两次报错均为PyTorch生态包与TPU专用torch_xla版本不兼容导致：

• 第一次报错是Colab默认预装的PyTorch 1.10与安装的torch_xla 1.9版本不匹配
• 第二次报错是降级PyTorch到1.9后，默认版本的torchtext与PyTorch 1.9不兼容

操作步骤 #

1. 首先重置Colab运行时，清除所有已安装的旧版本依赖：点击菜单栏「运行时」→「断开连接并删除运行时」，重新连接TPU实例。
2. 按如下顺序执行安装命令，所有依赖版本严格匹配PyTorch 1.9：

# 1. 安装TPU客户端和对应PyTorch 1.9的torch_xla
!pip install cloud-tpu-client==0.10 https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-1.9-cp37-cp37m-linux_x86_64.whl
# 2. 安装匹配PyTorch 1.9的生态包
!pip install torch==1.9.0 torchvision==0.10.0 torchaudio==0.9.0 torchtext==0.10.0
# 3. 安装兼容PyTorch 1.9的PyTorch Lightning版本
!pip install pytorch-lightning==1.5.10
# 4. 安装项目依赖，忽略torch相关包防止版本被覆盖
!pip install -r requirements.txt --ignore-installed torch torchvision torchaudio torchtext pytorch-lightning

3. 按照系统提示重启运行时，重启后不要重复执行上述安装命令，直接运行项目代码即可。

可选优化 #

如果你的项目本身不需要用到torchtext，可以直接卸载torchtext避免兼容问题：

!pip uninstall -y torchtext

PyTorch Lightning仅会检测torchtext是否可用，卸载后不会影响核心训练功能。

版本验证（可选） #

重启后可执行如下代码确认版本匹配：

import torch
print(torch.__version__) # 输出应为1.9.0
import torch_xla.core.xla_model as xm
import pytorch_lightning as pl
# 无报错即为版本匹配成功

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