torch.device用法总结
这里写自定义目录标题1 通过字符串构造设备2 通过字符串+序号构造设备3注意事项3.1可直接使用字符串构建3.2 有cuda的也可直接输入序号3.3 以下示例均等效torch.device代表将torch.Tensor分配到的设备的对象。torch.device包含一个设备类型('cpu'或'cuda'设备类型)和可选的设备的序号。如果设备序号不存在,则为当前设备;例如,torch.Tensor用
·
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
更多推荐
已为社区贡献1条内容
所有评论(0)