CUDA 上的张量
torch.device() 包含一个设备类型(‘cpu’ 或 ‘cuda’)和可选的设备序号,如果设备序号不存在,则为当前设备。
>>> torch.device('cpu')
device(type='cpu')
>>> torch.device('cuda')
device(type='cuda')
>>> torch.device('cpu', 0)
device(type='cpu', index=0)
>>> torch.device('cuda', 0)
device(type='cuda', index=0)
device_1 = torch.device('cpu')
device_2 = torch.device('cuda')
print(device_1.type)
print(device_2.type)
---------
cpu
cuda
张量可以使用 to() 方法移动到任何设备(device)上。
- tensor.to(‘cpu’)
- tensor.to(‘cuda’)
- device = torch.device(‘cuda’) --> tensor.to(device)
# 使用 torch.device() 将 tensor 移入和移出 GPU
x = torch.randn(1)
if torch.cuda.is_available():
device = torch.device('cuda') # CUDA 设备对象
y = torch.ones_like(x, device=device) # 直接在 GPU 上创建 tensor
x = x.to(device) # 也可以使用 x = x.to('cuda')
z = x + y
print(z)
print(z.to('cpu', torch.double)) # 张量在转移设备时可以改变 dtype
---------
tensor([1.0445], device='cuda:0')
tensor([1.0445], dtype=torch.float64)