查看gpu
python
import torch
# 检查CUDA是否可用
if torch.cuda.is_available():
# 获取GPU设备数量
device_count = torch.cuda.device_count()
# 列出可用的GPU设备
for i in range(device_count):
print(f"GPU {i}: {torch.cuda.get_device_name(i)}")
else:
print("CUDA is not available. No GPU devices found.")
输出结果
设置指定gpu运行
指定上面输出的5号卡
python
if torch.cuda.is_available():
# 指定要使用的GPU设备编号
device = torch.device("cuda:5")
print(f"Using GPU {device} - {torch.cuda.get_device_name(device)}")
else:
print("CUDA is not available. No GPU devices found.")
查看gpu的内存情况等
python
import subprocess
# 执行nvidia-smi命令以获取GPU信息
nvidia_smi_output = subprocess.check_output("nvidia-smi", shell=True).decode()
# 切分输出为每个GPU的信息
gpu_info = nvidia_smi_output.strip().split('\n\n')
# 遍历每个GPU的信息
for i, info in enumerate(gpu_info):
print(f"GPU {i}:")
print(info)