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

torch.device包含一个设备类型('cpu''cuda'设备类型)和可选的设备的序号。如果设备序号不存在,则为当前设备;
例如,torch.Tensor用设备构建'cuda'的结果等同于'cuda:X',其中Xtorch.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

Logo

为开发者提供学习成长、分享交流、生态实践、资源工具等服务,帮助开发者快速成长。

更多推荐