当前位置: 首页 > 工具软件 > Device > 使用案例 >

Pytorch 学习笔记--to(device)的用法

陆沈浪
2023-12-01

在学习深度学习的时候,我们写代码经常会见到类似的代码:

img = img.to(device=torch.device("cuda" if torch.cuda.is_available() else "cpu"))
model = models.vgg16_bn(pretrained=True).to(device=torch.device("cuda" if torch.cuda.is_available() else "cpu"))

也可以先定义device:

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

1. 那么段代码到底有什么用呢?

        这段代码的意思就是将所有最开始读取数据时的tensor变量copy一份到device所指定的GPU上去,之后的运算都在GPU上进行。

2. 你可能会问,为什么要在GPU上做运算呢?

        首先,在做高维特征运算的时候,采用GPU无疑是比用CPU效率更高,如果两个数据中一个加了.cuda()或者.to(device),而另外一个没有加,就会造成类型不匹配而报错。

        tensor和numpy都是矩阵,前者能在GPU上运行,后者只能在CPU运行,所以要注意数据类型的转换。

3. .cuda()和.to(device)的效果一样吗?为什么后者更好?

        两个方法都可以达到同样的效果,在pytorch中,即使是有GPU的机器,它也不会自动使用GPU,而是需要在程序中显示指定。调用model.cuda(),可以将模型加载到GPU上去。这种方法不被提倡,而建议使用model.to(device)的方式,这样可以显示指定需要使用的计算资源,特别是有多个GPU的情况下。

4. 如果你有多个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)

#本文为本人学习笔记,以上有部分内容转载自网络,若有侵权,请联系删除

 类似资料: