把数据迁移到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

 

这两个方法耗时基本差不多。