如何让Jupyter Notebook在多GPU环境下指定非0号GPU运行

为什么os.environ['CUDA_VISIBLE_DEVICES'] = '1'无效？ #

最常见的原因是你在导入PyTorch之后才设置这个环境变量——PyTorch在初始化时会读取CUDA环境变量，后续修改不会影响已经初始化的CUDA上下文。

有效解决方案 #

1. 提前设置环境变量（必须在导入torch之前） #

确保设置环境变量的代码是Notebook中最先执行的部分，顺序不能错：

import os
# 必须在import torch之前执行
os.environ['CUDA_VISIBLE_DEVICES'] = '1'

import torch
# 验证是否生效
print(f"可用GPU数量: {torch.cuda.device_count()}")
print(f"当前使用的GPU索引: {torch.cuda.current_device()}")

执行后如果输出当前使用的GPU索引: 0是正常的——因为CUDA_VISIBLE_DEVICES='1'会把物理1号GPU映射为逻辑0号，PyTorch内部只会看到这一个GPU。

2. 显式指定设备（无需依赖环境变量） #

直接在代码中为模型和张量指定目标GPU，这种方法更灵活，不受环境变量顺序影响：

import torch
# 指定使用物理1号GPU
device = torch.device('cuda:1')

# 将模型移到指定GPU
model = YourModel()
model = model.to(device)

# 将张量移到指定GPU
input_tensor = torch.randn(32, 3, 224, 224).to(device)
# 或者用cuda()方法指定索引
input_tensor = torch.randn(32, 3, 224, 224).cuda(1)

# 验证当前张量所在设备
print(input_tensor.device)

3. 启动Jupyter时指定GPU #

在终端启动Notebook时直接绑定GPU，这样整个Notebook进程都会默认使用指定GPU：

CUDA_VISIBLE_DEVICES=1 jupyter notebook

启动后，Notebook中所有PyTorch代码都会默认使用这个GPU，无需在代码中额外设置。

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