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-07 14:30
复旦大学 Java
遇到这种人我也不知道说啥了
无能的丈夫:但我觉得这个hr语气没什么问题啊(没有恶意
点赞 评论 收藏
分享
06-15 18:44
黄淮学院 Java
Lynn012:如果是居民楼还是算了吧,看着有点野呢
点赞 评论 收藏
分享
06-20 17:42
东华大学 Java
凉风落木楚山秋:要是在2015,你这简历还可以月入十万,可惜现在是2025,已经跟不上版本了
我的简历长这样
点赞 评论 收藏
分享
评论
点赞
收藏
分享

创作者周榜

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