torch.device用法总结

torch.device代表将torch.Tensor分配到的设备的对象。

torch.device包含一个设备类型（'cpu'或'cuda'设备类型）和可选的设备的序号。如果设备序号不存在，则为当前设备;
例如，torch.Tensor用设备构建'cuda'的结果等同于'cuda:X',其中X是torch.cuda.current_device()的结果。

torch.Tensor的设备可以通过Tensor.device访问属性。

构造torch.device可以通过字符串/字符串和设备编号。如下：

1 通过字符串构造设备

>>> torch.device('cpu')
device(type='cpu')

>>> torch.device('cuda')  # current cuda device
device(type='cuda')

2 通过字符串＋序号构造设备

>>> torch.device('cuda', 0)
device(type='cuda', index=0)

>>> torch.device('cpu', 0)
device(type='cpu', index=0)

3注意事项

3.1可直接使用字符串构建

torch.device函数中的参数通常可以用一个字符串替代。这允许使用代码快速构建原型。

>> # Example of a function that takes in a torch.device
>> cuda1 = torch.device('cuda:1')
>> torch.randn((2,3), device=cuda1)

>> # You can substitute the torch.device with a string
>> torch.randn((2,3), 'cuda:1')

两种方法等同

3.2 有cuda的也可直接输入序号

出于传统原因，可以通过单个设备序号构建设备，将其视为cuda设备。这匹配Tensor.get_device()，它为cuda张量返回一个序数，并且不支持cpu张量。

>> torch.device(1)
device(type='cuda', index=1)

3.3 以下示例均等效

指定设备的方法可以使用（properly formatted）字符串或（legacy）整数型设备序数，即以下示例均等效：

>> torch.randn((2,3), device=torch.device('cuda:1'))
>> torch.randn((2,3), device='cuda:1')
>> torch.randn((2,3), device=1)  # legacy

转自https://ptorch.com/news/187.html