torch.device('cuda') 与 torch.device('cuda:0')在进行计算时,对于单卡计算机而言,没有任何区别,都是唯一的那一张GPU。

示例如下:

>>> a = torch.device('cuda')
>>> b = torch.device('cuda:0')
>>> tensor_a = torch.ones(1, device=a)
>>> tensor_a
tensor([1.], device='cuda:0')
>>> tensor_b = torch.ones(1, device=b)
>>> tensor_b
tensor([1.], device='cuda:0')

通过以上示例可以发现,即使没有显式指定a使用gpu 0,但是pytorch还是将tensor分配在gpu 0上。

但是如果直接比较二者

>>> a == b
False

返回值居然是false,也就是二者表示的device并不相同,进一步观察a和b,

>>> a
device(type='cuda')
>>> b
device(type='cuda', index=0)

二者还是具有细微差别的,a的index为none,b的index为0,推测是这一点差别造成了二者比较返回false的原因。

所以如果需要判断,tensor是不是在某一个设备上,不能比较

>>> tensor_b.device == a
False
>>> tensor_a.device == a
False

而应该比较

>>> tensor_b.device == b
True
>>> tensor_a.device == b
True

通过以上比较可以发现,此时tensor_a与tensor_b的device均为torch.device('cuda:0'),因此在确定tensor是不是在某个设备上时,需要明确地指定设备的编号。

Logo

旨在为数千万中国开发者提供一个无缝且高效的云端环境,以支持学习、使用和贡献开源项目。

更多推荐