torch.device(‘cuda‘) 与 torch.device(‘cuda:0‘) 的区别简析

Posted 2021-05-17 mazinkaiser1991

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是不是在某个设备上时，需要明确地指定设备的编号。