如何在Python3中设置CUDA_VISIBLE_DEVICES指定GPU？

排查环境变量传递问题 #

• 先验证Python3进程是否读取到CUDA_VISIBLE_DEVICES变量，在代码开头添加：

  import os
  print(os.environ.get('CUDA_VISIBLE_DEVICES'))
  

  执行CUDA_VISIBLE_DEVICES=1 python3 program.py，若输出为None，说明环境变量未传递，可尝试用Python3绝对路径调用，比如：

  CUDA_VISIBLE_DEVICES=1 /usr/bin/python3 program.py
  

  避免alias或虚拟环境导致的环境变量隔离。

代码内提前设置环境变量 #

• 必须在导入任何CUDA相关库（如torch、tensorflow）之前设置环境变量，否则CUDA库初始化后读取的设备信息不会更新：

  import os
  os.environ['CUDA_VISIBLE_DEVICES'] = '1'
  # 之后再导入CUDA依赖库
  import torch
  

检查CUDA与Python3兼容性 #

• 确认Python3对应的CUDA Toolkit版本匹配，部分旧版CUDA对Python3的环境变量支持存在问题，可更新至适配的CUDA版本。

直接通过框架API指定设备 #

• 若上述方法无效，可直接调用CUDA框架的设备指定接口：
  • PyTorch示例：

    import torch
    # 全局指定设备
    torch.cuda.set_device(1)
    # 或使用上下文管理器临时指定
    with torch.cuda.device(1):
        # 执行你的代码逻辑
    

  • TensorFlow示例：

    import tensorflow as tf
    gpus = tf.config.list_physical_devices('GPU')
    tf.config.set_visible_devices(gpus[1], 'GPU')
    

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