把数据从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
这两个方法耗时基本差不多。