转载请注明出处：小锋学长生活大爆炸[xfxuezhagn.cn]
如果本文帮助到了你，欢迎[点赞、收藏、关注]哦~
问题示例
for gpus in ['0', '1', '2', '3', '4', '5', '6', '7']:os.environ['CUDA_VISIBLE_DEVICES'] = gpusprint(torch.cuda.get_device_name(0))始终将使用第一个GPU，即CUDA:0.
问题解析
在使用 os.environ['CUDA_VISIBLE_DEVICES'] 动态设置可见的 GPU 时，通常只能在程序开始运行之前进行设置，并且这种设置在程序运行后无法在一个进程中动态改变。这是因为深度学习框架（如 TensorFlow、PyTorch 等）在初始化时就会读取 CUDA_VISIBLE_DEVICES 环境变量，并在后续操作中使用这些设置。因此，尝试在运行时通过循环动态更改 CUDA_VISIBLE_DEVICES 是无效的。
解决方法
1. 子进程方法
通过在子进程中运行深度学习代码，每次子进程启动时重新设置 CUDA_VISIBLE_DEVICES 环境变量。在这种方法中，每次循环都会启动一个新的子进程 your_script.py，并为该子进程设置不同的 GPU。
import os
import subprocessgpus = ['0', '1', '2']for gpu in gpus:env = os.environ.copy()env['CUDA_VISIBLE_DEVICES'] = gpusubprocess.run(['python', 'your_script.py'], env=env)
2. 使用PyTorch的设备管理 (个人推荐)
在 PyTorch 中，你可以直接通过 torch.cuda.set_device(device) 动态设置设备，而无需更改 CUDA_VISIBLE_DEVICES。这种方法在单个进程中动态设置不同的 GPU。示例如下：
import torchgpus = [0, 1, 2]for gpu in gpus:torch.cuda.set_device(gpu)print(torch.cuda.get_device_name(torch.cuda.current_device()))# 你的深度学习代码model = your_model.to(torch.device(f'cuda:{gpu}'))# 继续进行训练或推理
3. 使用多进程方法
使用 multiprocessing 模块启动多个进程，每个进程设置不同的 GPU。这种方法启动多个独立进程，每个进程可以独立设置 CUDA_VISIBLE_DEVICES 并在其上运行代码。：
import os
from multiprocessing import Processdef run_on_gpu(gpu):os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu)# 你的深度学习代码# 如训练、推理等gpus = [0, 1, 2]
processes = []for gpu in gpus:p = Process(target=run_on_gpu, args=(gpu,))p.start()processes.append(p)for p in processes:p.join()
