【PyTorch】cuda()与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；