Pytorch模型数据的gpu和cpu:model.to(device), model.cuda()

背景介绍
我们在使用Pytorch训练时,模型和数据有可能加载在不同的设备上(gpu和cpu),在算梯度或者loss的时候,报错信息类似如下:

RuntimeError: Function AddBackward0 returned an invalid gradient at index 1 - expected type torch.cuda.FloatTensor but got torch.FloatTensor

RuntimeError: Expected object of device type cuda but got device type cpu for argument #2 'target' in call to _thnn_l1_loss_forward

这种报错通常有几种情况:

  1. 数据在cpu上,模型在gpu上;

  2. 数据在gpu上,模型在cpu上;

  3. 指定层不支持gpu或cpu版本,但数据和模型都在gpu或cpu上。

解决方案
第1步:将模型和数据都移到同一设备上。

假设模型和数据定义如下:

model = Model()

input = dataloader()

output = model(input)

移动到gpu上:

# solution: 0

device = 'gpu'

model = model.to(device)

data = data.to(device)


# solution: 1

model = model.cuda()

data = data.cuda()

移动到cpu上:

# solution: 0

device = 'cpu'

model = model.to(device)

data = data.to(device)


# solution: 1

model = model.cpu()

data = data.cpu()

第二步:打印模型model和数据data在gpu上还是cpu上。

  1. 通过判断模型model的参数是否在cuda上来判定模型是否在gpu上。
print('Is model on gpu: ', next(model.parameters()).is_cuda)
输出若是True,则model在gpu上;若是False,则model在cpu上。
  1. 输出数据data的device字段。
print('data device: ', data.device)
输出gpu则在gpu上,输出cpu则在cpu上。
全部评论

相关推荐

不愿透露姓名的神秘牛友
07-09 11:30
找工作7个月,投了7000封,3段世界五百强实习,才有一个offer,牛油们肯定比我强吧
码农索隆:不对不对不对,实习经历这么厉害,简历也没少投,问题出在哪呢
点赞 评论 收藏
分享
下个早班:秒挂就是不缺人
点赞 评论 收藏
分享
程序员小白条:这比例牛逼,750:1
点赞 评论 收藏
分享
评论
点赞
收藏
分享

创作者周榜

更多
牛客网
牛客网在线编程
牛客网题解
牛客企业服务