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