1. 检查GPU使用情况
由于nvidia-smi命令显示的gpu id和实际运行时gpu id并不对应,因此通过上述代码来查看实际运行时的gpu id
import torch
def print_gpu_info():
# 检查可用的 GPU 数量
num_gpus = torch.cuda.device_count()
if num_gpus > 0:
print(f"可用的 GPU 数量: {num_gpus}")
# 打印每个 GPU 设备的 ID 和名称
for device_id in range(num_gpus):
device_name = torch.cuda.get_device_name(device_id)
print(f"GPU 设备 {device_id}: {device_name}")
else:
print("没有可用的 GPU")
# 调用函数打印 GPU 信息
print_gpu_info()
# 设置要使用的 GPU 设备 ID
device_id = 0
torch.cuda.set_device(device_id)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Device:", device)
# 获取当前设备的 ID 和名称
device_id = torch.cuda.current_device()
device_name = torch.cuda.get_device_name(device_id)
# 打印使用的 GPU 设备
print(f"当前使用的 GPU 设备: {device_name}")
print('GPU id: ' + str(device_id))
2. 设置使用特定的GPU
在PyTorch中,你可以通过设置 CUDA_VISIBLE_DEVICES
环境变量来控制哪些GPU对当前进程可见,从而限制PyTorch使用的GPU数量。这个环境变量接受一个由逗号分隔的GPU索引列表。
如果你只想让PyTorch使用特定的几个GPU,可以在运行Python脚本或程序之前设置这个环境变量。
在命令行中,你可以这样设置(假设你想使用GPU 0和GPU 2):
export CUDA_VISIBLE_DEVICES=0,2
然后你可以运行你的PyTorch程序,程序将只看到你指定的GPU。
在Python脚本中,你可以使用 os.environ
在程序开始时设置:
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,2"
import torch
# 现在PyTorch只能看到GPU 0和GPU 2
请注意,当你设置 CUDA_VISIBLE_DEVICES
后,PyTorch看到的GPU索引将会重新排序以匹配这个变量。也就是说,即使你设置 CUDA_VISIBLE_DEVICES=2,3
,在PyTorch中这两个设备仍然会被标记为 cuda:0
和 cuda:1
。