PyTorch多GPU训练触发Invalid device id错误求助

问题场景 #

基于pytorch-deeplab-xception仓库开发DeepLab V3模型时，已在程序中设置os.environ["CUDA_VISIBLE_DEVICES"] = "1,2,3"，但执行以下代码时触发AssertionError: Invalid device id：

self.model = torch.nn.DataParallel(self.model, device_ids=[1,2,3])

完整错误栈：

Traceback (most recent call last):
  File "train.py", line 313, in <module>
    main()
  File "train.py", line 302, in main
    trainer = Trainer(args)
  File "train.py", line 76, in __init__
    self.model = torch.nn.DataParallel(self.model, device_ids=[1,2,3]).to(args.cuda)
  File "/home/deshpand/anaconda3/envs/torch_env/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py", line 145, in __init__
    _check_balance(self.device_ids)
  File "/home/deshpand/anaconda3/envs/torch_env/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py", line 25, in _check_balance
    dev_props = _get_devices_properties(device_ids)
  File "/home/deshpand/anaconda3/envs/torch_env/lib/python3.8/site-packages/torch/_utils.py", line 577, in _get_devices_properties
    return [_get_device_attr(lambda m: m.get_device_properties(i)) for i in device_ids]
  File "/home/deshpand/anaconda3/envs/torch_env/lib/python3.8/site-packages/torch/_utils.py", line 577, in <listcomp>
    return [_get_device_attr(lambda m: m.get_device_properties(i)) for i in device_ids]
  File "/home/deshpand/anaconda3/envs/torch_env/lib/python3.8/site-packages/torch/_utils.py", line 558, in _get_device_attr
    return get_member(torch.cuda)
  File "/home/deshpand/anaconda3/envs/torch_env/lib/python3.8/site-packages/torch/_utils.py", line 577, in <lambda>
    return [_get_device_attr(lambda m: m.get_device_properties(i)) for i in device_ids]
  File "/home/deshpand/anaconda3/envs/torch_env/lib/python3.8/site-packages/torch/cuda/__init__.py", line 374, in get_device_properties
    raise AssertionError("Invalid device id")
AssertionError: Invalid device id

解决方案 #

1. 修正device_ids的逻辑ID
   当设置CUDA_VISIBLE_DEVICES后，PyTorch会将指定的物理GPU映射为逻辑上的连续ID（从0开始）。此时你指定的[1,2,3]是物理设备ID，但PyTorch只能识别映射后的逻辑ID[0,1,2]。修改代码为：

# 显式指定映射后的逻辑GPU ID
self.model = torch.nn.DataParallel(self.model, device_ids=[0,1,2]).to(args.cuda)
# 或省略device_ids参数，自动使用所有可见GPU
self.model = torch.nn.DataParallel(self.model).to(args.cuda)

2. 确保环境变量设置时机正确
   os.environ["CUDA_VISIBLE_DEVICES"]必须在导入torch之前设置，否则不会生效。调整代码顺序为：

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1,2,3"
import torch
# 后续模型初始化、训练代码

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