不乱于心,不困于情。
不畏将来,不念过往。如此,安好。

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

本文于 2024-07-09 17:30 更新,部分内容具有时效性,如有失效,请留言

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程序
CUDA_VISIBLE_DEVICES=0 python extract_masks.py

方式四(多卡):命令行里,在GPU的id为0和1的两张显卡上运行**.py程序:
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

赞(0) 打赏
未经允许不得转载:seo优化_前端开发_渗透技术 » python代码查看可用的gpu,然后指定可用的gpu运行

相关推荐

  • 暂无文章

觉得文章有用就打赏一下文章作者

支付宝扫一扫打赏

微信扫一扫打赏