在Google Colab Pro中执行Python脚本无法调用GPU的求助

可能的原因及对应方案 #

1. 脚本main.py未指定使用CUDA #

Notebook能检测到GPU，但脚本本身可能未将模型和数据迁移到CUDA设备。需要在main.py中添加设备配置逻辑：

import torch

# 检测并设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 将模型移至GPU
model = YourModelClass().to(device)

# 训练/推理时同步数据到GPU
for data in dataloader:
    inputs, labels = data
    inputs, labels = inputs.to(device), labels.to(device)
    # 后续训练逻辑

2. 确保环境变量与Notebook共享 #

用!python执行脚本时，可能未继承Notebook的CUDA环境配置。可以：

• 改用Colab魔法命令%run main.py运行脚本，它会直接共享Notebook的运行环境，包括CUDA配置。
• 若坚持用!python，显式指定CUDA设备：

!CUDA_VISIBLE_DEVICES=0 python main.py

3. 检查PyTorch版本一致性 #

Notebook和脚本使用的PyTorch版本不一致，可能导致CUDA检测失败。分别在Notebook和脚本中运行以下代码确认：

import torch
print(torch.__version__)
print(torch.cuda.is_available())

若版本不符，在Notebook中重新安装对应CUDA版本的PyTorch，例如：

!pip install torch==2.0.1+cu118 torchvision==0.15.2+cu118 --index-url https://download.pytorch.org/whl/cu118

4. 验证GPU资源状态 #

Colab的GPU可能被其他进程占用，导致脚本无法调用。在Notebook中运行以下命令查看GPU使用情况：

!nvidia-smi

若显示GPU被占用，可重启Colab运行时（菜单栏「运行时」→「重启运行时」）释放资源。

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