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

小怪兽会微笑

已于 2023-11-25 09:56:56 修改

阅读量5.2k

点赞数 1

分类专栏：学习笔记文章标签： python 深度学习 pytorch

于 2023-10-31 11:04:59 首次发布

本文链接：https://blog.csdn.net/a1920993165/article/details/134135823

版权

学习笔记专栏收录该内容

140 篇文章

订阅专栏

文章介绍了如何在Python中使用nvidia-smi查看GPU信息，设置单卡或多卡环境，以及如何指定特定GPU运行模型。方法包括检查CUDA可用性，设置CUDA_VISIBLE_DEVICES环境变量，以及在多GPU环境下的训练和测试操作。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

1.查看gpu

（最好像下面这样代码里输出看下，nvidia-smi里gpu编号有时候对不上的）

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.")

输出结果
在这里插入图片描述
python可以借助nvidia-smi查看gpu的内存情况等

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)

2.设置指定gpu运行

方式一（单卡）：设置可见的为1号卡，再指定0号卡运行（设置可见卡为1号后，默认为0号卡开始编号）

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

if torch.cuda.is_available():
    # 指定要使用的GPU设备编号
    device = torch.device("cuda:0")
    print(f"Using GPU {device} - {torch.cuda.get_device_name(device)}")
else:
    print("CUDA is not available. No GPU devices found.")

方式二（单卡）：直接指定5号卡

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的id为0的两张显卡上运行**.py程序

export CUDA_VISIBLE_DEVICES=0 
python extract_masks.py

方式四（多卡）：命令行里，在GPU的id为0和1的两张显卡上运行**.py程序：

export CUDA_VISIBLE_DEVICES=0,1
python ***.py

方式五（多卡）：指定可见的为0，2，3号卡

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,2,3"
# 注意：这两行代码必须在文件的最开头，在加载各种包之前

model = ......
model.train(True) or model.train(False) # 看你是要训练还是测试
model.to('cuda') # 或者model.cuda()

方式六（多卡）：直接代码指定0，1号卡

import torch

# 检查CUDA是否可用
if torch.cuda.is_available():
    # 指定要使用的GPU设备编号，可以使用逗号分隔多个GPU编号
    device = torch.device("cuda:0,1")
    print(f"Using GPUs {device} - {torch.cuda.get_device_name(0)} and {torch.cuda.get_device_name(1)}")
else:
    print("CUDA is not available. No GPU devices found.")

# 在后续的代码中，使用 device 参数来指定运行在特定的GPU上
# 例如1：
model = model.to(device)

# 例如2：
import torch
tensor = torch.randn(3, 3)# 创建一个张量
tensor = tensor.to(device)# 将张量移动到指定的GPU