【Pytorch】指定训练时使用的GPU ID

1. 在.py文件中添加语句：os.environ['CUDA_VISIBLE_DEVICES'] = '3,4,5'

这里的数字并不代表GPU对应的ID，相反，这里的数字有{ }个，就只能使用ID为前{ }的device。例如上例使用的就是ID为0、1、2的device。
同理：
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2'
os.environ['CUDA_VISIBLE_DEVICES'] = '5,9,10'
效果都是相同的。

2. 在命令行运行时：CUDA_VISIBLE_DEVICES=1,2 python myscript.py
   这里的1,2指的就是nvidia-smi看到的实际的device id。

ref: https://discuss.pytorch.org/t/how-to-change-the-default-device-of-gpu-device-ids-0/1041

3. 在.py文件中添加：
   此时设置的cuda:0为GPU对应的ID

import torch

# 检测是否有可用的 CUDA 设备
if torch.cuda.is_available():
    device = torch.device('cuda:0')  # 将设备设置为第一个 CUDA 设备
    print('Using GPU:', torch.cuda.get_device_name(device))
else:
    device = torch.device('cpu')  # 如果没有可用的 CUDA 设备，就使用 CPU
    print('Using CPU')

# 后面的代码可以使用 `device` 变量将 tensors 或模型发送到正确的设备 
# 例如，你可以使用 `tensor.to(device)` 或者 `model.to(device)`