python指定Gpu

1、目前主流方法:.to(device)方法(推荐)

import torch
import time

#1.通常用法
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
data = data.to(device)
model = model.to(device)
'''
    1.先创建device对象
    2.to函数指定数据或者模型放到哪里
'''

#2.将构建的tensor或者模型放到指定设备上(GPU)
torch.device('cuda',0) #这里的0指的是设备的序号
torch.device('cuda:0')

#3.例子 cpu转到GPU
s = time.time()
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
a = torch.rand([1024, 1024, 10]).to(device)
print('time:{:6.3f}'.format(time.time()-s)) # 输出: time: 0.087

#4.例子2 Tensor GPU 转到cpu运行

 predict_boxes = predictions["boxes"].to("cpu").numpy()
 predict_classes = predictions["labels"].to("cpu").numpy()
 predict_scores = predictions["scores"].to("cpu").numpy()

.cuda()方法和.to(device)方法耗时基本差不多。但是 .to(device)代码风格更容易后期修改。

2、to(device)需要注意的是:

使用GPU训练的时候,需要将Module对象和Tensor类型的数据送入到device。通常会使用 to.(device)。但是需要注意的是:

  • 对于Tensor类型的数据,使用to.(device) 之后,需要接收返回值,返回值才是正确设置了device的Tensor。
  • 对于Module对象,只用调用to.(device) 就可以将模型设置为指定的device。不必接收返回值。
Module对象设置device的写法
model.to(device)

Tensor类型的数据设置 device 的写法。
samples = samples.to(device)

当然model也可以接收返回值。

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)
img = img.to(device)

官网解释:

python指定Gpu

python指定Gpu

3、如果你有多个GPU

这样可以显示指定需要使用的计算资源,特别是有多个GPU的情况下。

那么可以参考以下代码:

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = Model()

if torch.cuda.device_count() > 1:
    model = nn.DataParallel(model,device_ids=[0,1,2])
    model.to(device)

1、torch.device用法总结_会飞的萨摩的博客-CSDN博客主要参考这两篇文章:

1、torch.device用法总结_会飞的萨摩的博客-CSDN博客

2、pytorch并行处理详解(多GPU,环境变量)_xys430381_1的博客-CSDN博客

Original: https://blog.csdn.net/qimo601/article/details/123822178
Author: 三世
Title: python指定Gpu

原创文章受到原创版权保护。转载请注明出处:https://www.johngo689.com/672855/

转载文章受原作者版权保护。转载请注明原作者出处!

(0)

大家都在看

亲爱的 Coder【最近整理,可免费获取】👉 最新必读书单  | 👏 面试题下载  | 🌎 免费的AI知识星球