Pytorch文档——Tensor Attributes -torch.device

torch.device
class torch.device
A一个torch.device 是一个对象，这个对象代表着torch.tensor在哪个device上，或者会被分配到哪个device。

torch.device 对象包含 device type (‘cpu’ or ‘cuda’) 还有一个可选择的设备类型序列号. 如果设备序列号没有被声明，这代表着使用当前的设备 。例如一个 torch.Tensor 使用 'cuda’作为运行设备，这与 使用 ‘cuda:X’ X 是 torch.cuda.current_device()的结果一样.

torch.Tensor的设备号可以通过 Tensor.device获得.

torch.device 可以通过一个字符串或者通过带有序列号的字符串指定

Via a string:

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

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

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

Via a string and device ordinal:

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

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

The torch.device argument in functions can generally be substituted with a string. This allows for fast prototyping of code.

>>> # 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')

Note

For legacy reasons, a device can be constructed via a single device ordinal, which is treated as a cuda device. This matches Tensor.get_device(), which returns an ordinal for cuda tensors and is not supported for cpu tensors.

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

Note

Methods which take a device will generally accept a (properly formatted) string or (legacy) integer device ordinal, i.e. the following are all equivalent:

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