怎么用python代码查看可用的gpu,然后指定可用的gpu运行

查看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)
相关推荐
IMPYLH16 分钟前
Python 的内置函数 reversed
笔记·python
shangyingying_12 小时前
关于小波降噪、小波增强、小波去雾的原理区分
人工智能·深度学习·计算机视觉
小赖同学啊2 小时前
物联网数据安全区块链服务
开发语言·python·区块链
码荼3 小时前
学习开发之hashmap
java·python·学习·哈希算法·个人开发·小白学开发·不花钱不花时间crud
书玮嘎3 小时前
【WIP】【VLA&VLM——InternVL系列】
人工智能·深度学习
要努力啊啊啊3 小时前
YOLOv2 正负样本分配机制详解
人工智能·深度学习·yolo·计算机视觉·目标跟踪
小陈phd4 小时前
李宏毅机器学习笔记——梯度下降法
人工智能·python·机器学习
喝过期的拉菲4 小时前
如何使用 Pytorch Lightning 启用早停机制
pytorch·lightning·早停机制
kk爱闹4 小时前
【挑战14天学完python和pytorch】- day01
android·pytorch·python
Blossom.1184 小时前
机器学习在智能建筑中的应用:能源管理与环境优化
人工智能·python·深度学习·神经网络·机器学习·机器人·sklearn