pytorch-gpu的使用

2022-10-17 12:36:31

1. 服务器GPU状态查询

1) lspci | grep -i nvidia 可以查询所有nvidia显卡
2) lspci -v -s [显卡编号] 可以查看显卡具体属性
3) nvidia-smi 可以查看显卡的显存利用率

2. torch.cuda主要函数

1) 查看是否有可用GPU、可用GPU数量: torch.cuda.is_available(), torch.cuda.device_count()
2) 查看当前使用的GPU序号:torch.cuda.current_device()
3) 查看指定GPU的容量、名称:
    torch.cuda.get_device_capability(device),  torch.cuda.get_device_name(device)
4) 清空程序占用的GPU资源: torch.cuda.empty_cache()
5) 为GPU设置随机种子:torch.cuda.manual_seed(seed), torch.cuda.manual_seed_all(seed)

设置某gpu加速

os.environ["CUDA_VISIBLE_DEVICES"] =''X"

X为设置的gpu编号

os.environ["CUDA_VISIBLE_DEVICES"] =''1,2,3"的含义是指定gpu1号,2号,3号可见,
device = torch.device("cuda")

或者

device = torch.device("cuda:0")

指定使用0号gpu,同时在上面我们设置了1,2,3号gpu可见,因此此处的cuda:0就是1号gpu,

将模型放到cuda:0上

model.to(device)

将数据放到cuda:0上

image = image.to(device)

注:
model.to(device)可以直接将model的parameters放到gpu上,
但是image.to(device)不可以,必须进行赋值操作

总结:

os.environ["CUDA_VISIBLE_DEVICES"] =''1,2,3" # 含义是指定gpu1号,2号,3号可见,
device = torch.device("cuda") 
model.to(device)
image = image.to(device)

查看当前数据在哪个设备上

print(images.device)

查看当前模型是否使用gpu加速

print(next(model.parameters()).device)
  • 作者:RichardorMu
  • 原文链接:https://blog.csdn.net/weixin_42264234/article/details/118797827
    更新时间:2022-10-17 12:36:31