Python调用GPU的设备检测与选择：步骤详解与代码示例

核心步骤

1. 环境检测

   • 检查GPU驱动是否安装：通过nvidia-smi命令验证（需提前安装NVIDIA驱动）
   • Python库依赖：安装torch或tensorflow，推荐使用PyTorch

2. 设备检测逻辑

   • 优先级：GPU > CPU
   • 多GPU时默认选择索引0的设备，也可手动指定

3. 设备切换

   • 将张量（Tensor）或模型显式移动到目标设备

代码示例（PyTorch实现）

import torch

# 步骤1：检测可用设备
def detect_devices():
    if torch.cuda.is_available():
        num_gpus = torch.cuda.device_count()
        print(f"检测到 {num_gpus} 块GPU:")
        for i in range(num_gpus):
            print(f"  GPU {i}: {torch.cuda.get_device_name(i)}")
        return "cuda"  # 默认使用首块GPU
    else:
        print("无可用GPU，使用CPU")
        return "cpu"

# 步骤2：选择设备
device_type = detect_devices()
device = torch.device(device_type)

# 步骤3：使用设备（示例：创建张量并运行计算）
x = torch.tensor([1.0, 2.0, 3.0]).to(device)  # 显式移动到设备
y = torch.tensor([4.0, 5.0, 6.0]).to(device)
z = x + y
print(f"计算结果（设备: {device}）: {z}")

关键函数说明

多GPU手动选择示例

# 指定使用第二块GPU（索引从0开始）
if torch.cuda.device_count() >= 2:
    device = torch.device("cuda:1") 
    print(f"手动选择: GPU 1 ({torch.cuda.get_device_name(1)})")

注意事项

1. 设备一致性：所有交互的张量必须在同一设备，否则报错

   # 错误示例：x在GPU，y在CPU
   x = torch.tensor([1]).cuda()
   y = torch.tensor([2])
   z = x + y  # RuntimeError!
   

2. 内存管理：使用.cpu()将数据移回CPU处理大规模结果
3. 库差异：TensorFlow使用tf.config.list_physical_devices('GPU')检测设备

通过合理选择设备，GPU加速可使计算速度提升$10\times$至$100\times$（视任务复杂度而定）