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

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

原文地址:https://www.cnblogs.com/picassooo/p/13749666.html