CUDA_VISIBLE_DEVICES设置无效及PyTorch多GPU配置求助

先解决单设备指定无效的问题 #

先把基础问题修正，再推进多GPU配置：

• 命令行环境变量设置要严格按顺序：必须将CUDA_VISIBLE_DEVICES放在python命令前，且无多余空格，正确写法：

  CUDA_VISIBLE_DEVICES=1 python mytest.py
  

  可在代码开头加以下代码验证环境变量是否生效：

  import os
  print("CUDA_VISIBLE_DEVICES:", os.environ.get('CUDA_VISIBLE_DEVICES', '未设置'))
  

• 代码内设置环境变量必须在导入PyTorch之前：PyTorch启动时会读取CUDA设备信息，后续修改环境变量无效，正确顺序：

  import os
  os.environ['CUDA_VISIBLE_DEVICES'] = '1'  # 多设备写'0,1'
  import torch
  

• torch.device("cuda:1")需配合数据/模型迁移：仅定义设备不会自动切换，需将模型和张量移至该设备：

  device = torch.device("cuda:1")
  model = model.to(device)
  tensor = tensor.to(device)
  

多GPU配置的正确方法 #

torch.cuda.set_device()仅支持指定单个默认设备，传入多设备字符串必然报错，以下是两种常用多GPU配置方式：

方式1：环境变量+DataParallel（自动多GPU并行） #

1. 指定多可见设备：
   命令行方式：

   CUDA_VISIBLE_DEVICES=0,1 python mytest.py
   

   或代码开头（导入torch前）设置：

   import os
   os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
   import torch
   

2. 用DataParallel包装模型：

   model = torch.nn.DataParallel(model)
   

   该方式会自动将模型拆分到所有可见GPU，数据自动分片计算并合并结果，无需手动管理设备。

方式2：上下文管理器手动控制多设备（适合自定义逻辑） #

若需精细控制不同任务使用不同GPU，可使用设备上下文：

# 在GPU 0上执行任务
with torch.cuda.device('cuda:0'):
    tensor0 = torch.randn(10).cuda()
    model0 = model.to('cuda:0')
    output0 = model0(tensor0)

# 在GPU 1上执行任务
with torch.cuda.device('cuda:1'):
    tensor1 = torch.randn(10).cuda()
    model1 = model.to('cuda:1')
    output1 = model1(tensor1)

验证多GPU生效状态 #

加入以下代码确认设备识别情况：

print("可用GPU数量:", torch.cuda.device_count())
for i in range(torch.cuda.device_count()):
    print(f"GPU {i}: {torch.cuda.get_device_name(i)}")

若输出显示2个GPU，说明环境变量设置成功；使用DataParallel后，可通过print(next(model.parameters()).device)查看模型参数所在设备（会显示主设备，实际多GPU均有模型副本）。

额外注意事项 #

• 确保PyTorch版本与CUDA 10.2兼容（推荐PyTorch 1.6~1.9版本），版本不兼容可能导致设备识别异常。
• 避免混合使用多种设备指定方式（如同时用CUDA_VISIBLE_DEVICES和torch.cuda.set_device），防止出现逻辑冲突。

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