PyTorch中通过CUDA_VISIBLE_DEVICES指定GPU及排障

我写了以下代码尝试通过环境变量指定GPU，但一直报错：

def get_device_via_env_variables(deterministic: bool = False, verbose: bool = True) -> torch.device:
    device: torch.device = torch.device("cpu")
    if torch.cuda.is_available():
        if 'CUDA_VISIBLE_DEVICES' not in os.environ:
            device: torch.device = torch.device("cuda:0")
        else:
            gpu_idx: list[str] = os.environ['CUDA_VISIBLE_DEVICES'].split(',')
            if len(gpu_idx) == 1:
                gpu_idx: str = gpu_idx[0]
            else:
                # generate random int from 0 to len(gpu_idx) with import statement
                import random
                idx: int = random.randint(0, len(gpu_idx) - 1) if not deterministic else -1
                gpu_idx: str = gpu_idx[idx]
            device: torch.device = torch.device(f"cuda:{gpu_idx}")
    if verbose:
        print(f'{device=}')
    return device

我怀疑gpu_idx和CUDA_VISIBLE_DEVICES不匹配，希望能正确加载目标GPU。当前遇到两个错误：

1. 反序列化错误 #

Traceback (most recent call last):aded (0.000 MB deduped)
  File "/lfs/ampere1/0/brando9/diversity-for-predictive-success-of-meta-learning/div_src/diversity_src/experiment_mains/main_experiment_analysis_sl_vs_maml_performance_comp_distance.py", line 1368, in <module>
    main_data_analyis()
  File "/lfs/ampere1/0/brando9/diversity-for-predictive-success-of-meta-learning/div_src/diversity_src/experiment_mains/main_experiment_analysis_sl_vs_maml_performance_comp_distance.py", line 1163, in main_data_analyis
    args: Namespace = load_args()
  File "/lfs/ampere1/0/brando9/diversity-for-predictive-success-of-meta-learning/div_src/diversity_src/experiment_mains/main_experiment_analysis_sl_vs_maml_performance_comp_distance.py", line 1152, in load_args
    args.meta_learner = get_maml_meta_learner(args)
  File "/afs/cs.stanford.edu/u/brando9/diversity-for-predictive-success-of-meta-learning/div_src/diversity_src/data_analysis/common.py", line 272, in get_maml_meta_learner
    base_model = load_model_ckpt(args, path_to_checkpoint=args.path_2_init_maml)
  File "/afs/cs.stanford.edu/u/brando9/ultimate-utils/ultimate-utils-proj-src/uutils/torch_uu/mains/common.py", line 265, in load_model_ckpt
    base_model, _, _ = load_model_optimizer_scheduler_from_ckpt(args, path_to_checkpoint,
  File "/afs/cs.stanford.edu/u/brando9/ultimate-utils/ultimate-utils-proj-src/uutils/torch_uu/mains/common.py", line 81, in load_model_optimizer_scheduler_from_ckpt
    ckpt: dict = torch.load(path_to_checkpoint, map_location=torch.device('cuda:3'))
  File "/lfs/ampere1/0/brando9/miniconda/envs/mds_env_gpu/lib/python3.9/site-packages/torch/serialization.py", line 607, in load
    return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
  File "/lfs/ampere1/0/brando9/miniconda/envs/mds_env_gpu/lib/python3.9/site-packages/torch/serialization.py", line 882, in _load
    result = unpickler.load()
  File "/lfs/ampere1/0/brando9/miniconda/envs/mds_env_gpu/lib/python3.9/site-packages/torch/serialization.py", line 857, in persistent_load
    load_tensor(data_type, size, key, _maybe_decode_ascii(location))
  File "/lfs/ampere1/0/brando9/miniconda/envs/mds_env_gpu/lib/python3.9/site-packages/torch/serialization.py", line 846, in load_tensor
    loaded_storages[key] = restore_location(storage, location)
  File "/lfs/ampere1/0/brando9/miniconda/envs/mds_env_gpu/lib/python3.9/site-packages/torch/serialization.py", line 827, in restore_location
    return default_restore_location(storage, str(map_location))
  File "/lfs/ampere1/0/brando9/miniconda/envs/mds_env_gpu/lib/python3.9/site-packages/torch/serialization.py", line 175, in default_restore_location
    result = fn(storage, location)
  File "/lfs/ampere1/0/brando9/miniconda/envs/mds_env_gpu/lib/python3.9/site-packages/torch/serialization.py", line 151, in _cuda_deserialize
    device = validate_cuda_device(location)
  File "/lfs/ampere1/0/brando9/miniconda/envs/mds_env_gpu/lib/python3.9/site-packages/torch/serialization.py", line 142, in validate_cuda_device
    raise RuntimeError('Attempting to deserialize object on CUDA device '
RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.

2. CUDA内存不足错误（期望用GPU3但实际用GPU0） #

Traceback (most recent call last):
  File "/lfs/ampere1/0/brando9/diversity-for-predictive-success-of-meta-learning/div_src/diversity_src/experiment_mains/main_experiment_analysis_sl_vs_maml_performance_comp_distance.py", line 1368, in <module>
    main_data_analyis()
  File "/lfs/ampere1/0/brando9/diversity-for-predictive-success-of-meta-learning/div_src/diversity_src/experiment_mains/main_experiment_analysis_sl_vs_maml_performance_comp_distance.py", line 1213, in main_data_analyis
    stats_analysis_with_emphasis_on_effect_size(args, hist=True)
  File "/afs/cs.stanford.edu/u/brando9/diversity-for-predictive-success-of-meta-learning/div_src/diversity_src/data_analysis/stats_analysis_with_emphasis_on_effect_size.py", line 74, in stats_analysis_with_emphasis_on_effect_size
    results_usl: dict = get_episodic_accs_losses_all_splits_usl(args, args.mdl_sl, loaders)
  File "/afs/cs.stanford.edu/u/brando9/diversity-for-predictive-success-of-meta-learning/div_src/diversity_src/data_analysis/common.py", line 616, in get_episodic_accs_losses_all_splits_usl
    losses, accs = agent.get_lists_accs_losses(data, training)
  File "/afs/cs.stanford.edu/u/brando9/ultimate-utils/ultimate-utils-proj-src/uutils/torch_uu/meta_learners/pretrain_convergence.py", line 92, in get_lists_accs_losses
    spt_embeddings_t = self.get_embedding(spt_x_t, self.base_model).detach()
  File "/afs/cs.stanford.edu/u/brando9/ultimate-utils/ultimate-utils-proj-src/uutils/torch_uu/meta_learners/pretrain_convergence.py", line 166, in get_embedding
    return get_embedding(x=x, base_model=base_model)
  File "/afs/cs.stanford.edu/u/brando9/ultimate-utils/ultimate-utils-proj-src/uutils/torch_uu/meta_learners/pretrain_convergence.py", line 267, in get_embedding
    out = base_model.model.features(x)
  File "/lfs/ampere1/0/brando9/miniconda/envs/mds_env_gpu/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/lfs/ampere1/0/brando9/miniconda/envs/mds_env_gpu/lib/python3.9/site-packages/torch/nn/modules/container.py", line 139, in forward
    input = module(input)
  File "/lfs/ampere1/0/brando9/miniconda/envs/mds_env_gpu/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/lfs/ampere1/0/brando9/miniconda/envs/mds_env_gpu/lib/python3.9/site-packages/torch/nn/modules/conv.py", line 443, in forward
    return self._conv_forward(input, self.weight, self.bias)
  File "/lfs/ampere1/0/brando9/miniconda/envs/mds_env_gpu/lib/python3.9/site-packages/torch/nn/modules/conv.py", line 439, in _conv_forward
    return F.conv2d(input, weight, bias, self.stride,
RuntimeError: CUDA out of memory. Tried to allocate 174.00 MiB (GPU 0; 79.20 GiB total capacity; 54.31 GiB already allocated; 22.56 MiB free; 54.61 GiB reserved in total by PyTorch)

我希望使用GPU3，但最后一个错误显示用的是GPU0，请问哪里出错了？

问题分析与修复方案 #

核心问题1：CUDA_VISIBLE_DEVICES的逻辑误解

当你设置CUDA_VISIBLE_DEVICES=3时，PyTorch会把这个GPU映射为虚拟的cuda:0，而非保留原设备编号。你的代码直接用gpu_idx作为设备编号，比如原GPU3被映射后，代码里应该用cuda:0而非cuda:3，这直接导致两个错误：

• 反序列化时硬编码map_location=torch.device('cuda:3')，但此时PyTorch仅能看到1个虚拟GPU（cuda:0），因此报错设备不存在。
• 代码生成的cuda:{gpu_idx}指向系统原GPU编号，被CUDA_VISIBLE_DEVICES限制后PyTorch无法识别，最终默认回退到cuda:0，引发内存不足。

核心问题2：模型加载时硬编码设备

在load_model_optimizer_scheduler_from_ckpt函数中，你直接写死了map_location=torch.device('cuda:3')，完全忽略CUDA_VISIBLE_DEVICES的环境配置，必须改为动态获取当前设备。

修复后的设备获取代码

import os
import torch
import random

def get_device_via_env_variables(deterministic: bool = False, verbose: bool = True) -> torch.device:
    device: torch.device = torch.device("cpu")
    if torch.cuda.is_available():
        # 当设置CUDA_VISIBLE_DEVICES后，PyTorch仅可见指定GPU，且编号从0开始
        device_count = torch.cuda.device_count()
        if device_count == 0:
            device = torch.device("cpu")
        elif device_count == 1:
            device = torch.device("cuda:0")
        else:
            # 多GPU时随机选或固定选第一个
            if deterministic:
                device = torch.device("cuda:0")
            else:
                idx = random.randint(0, device_count - 1)
                device = torch.device(f"cuda:{idx}")
    if verbose:
        print(f'{device=}, 可见GPU数量: {torch.cuda.device_count()}')
        if device.type == 'cuda':
            print(f'当前使用GPU名称: {torch.cuda.get_device_name(device)}')
    return device

额外修复点：模型加载的map_location

把硬编码的map_location=torch.device('cuda:3')改为动态获取的设备：

# 先获取设备
device = get_device_via_env_variables()
# 加载模型时使用该设备
ckpt = torch.load(path_to_checkpoint, map_location=device)

验证步骤

1. 启动脚本前设置环境变量：export CUDA_VISIBLE_DEVICES=3
2. 运行代码，查看打印的device是否为cuda:0（此为虚拟编号，对应系统的GPU3）
3. 确认模型加载时使用了正确的设备，避免硬编码。

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