把數據遷移到GPU: .cuda()方法和.to(device)方法

把數據從CPU遷移到GPU時，可以用.cuda()方法，也可以用.to(device)方法。示例如下。

.cuda()方法

import torch
import time

t = time.time()
b = torch.rand([1024, 1024, 10])
b = b.cuda()
print('time:{:6.3f}'.format(time.time() - t))  # 輸出: time: 0.084

 

.to(device)方法

import torch
import time

s = time.time()
device = torch.device("cuda")
a = torch.rand([1024, 1024, 10]).to(device)
print('time:{:6.3f}'.format(time.time()-s))   # 輸出: time: 0.087

 

這兩個方法耗時基本差不多。