PyTorchcuda()与to(device)的区别

Posted 算法与编程之美

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了PyTorchcuda()与to(device)的区别相关的知识,希望对你有一定的参考价值。

问题

PyTorch中的Tensor要想在GPU中运行,可以有两种实现方式,其一是x.cuda(),其二是x.to(device)。两种方式均能实现GPU上运行,那么二者的区别是什么呢?

方法

import torch


device = 'cuda' if torch.cuda.is_available() else 'cpu'

a = torch.randn([3, 224, 224])

# (1) cuda()方式代码扩展性差,如果机器不支持GPU,则需要修改代码后才能在CPU上运行;
a.cuda()

# (2) to(device)方式代码扩展性好,即使没有GPU,也可以运行代码,不用做任何修改;
a.to(device)



结语

  • 推荐使用to(device)的方式,主要原因在于这样的编程方式更加易于扩展,而cuda()必须要求机器有GPU,否则需要修改所有代码,to(device)的方式则不受此限制,device既可以是CPU也可以是GPU;

以上是关于PyTorchcuda()与to(device)的区别的主要内容,如果未能解决你的问题,请参考以下文章

PyTorchCUDA Toolkit 及显卡驱动版本对应关系

Pytorch 设备和 .to(device) 方法

.to(device)和.cuda()设置GPU的区别

Attempting to fetch value instead of handling error Internal: failed to get device attribute 13 for

`R.string.cast_casting_to_device`在哪个包中

Pytorch Lighting 的 model.to(device)