pytorch 载入历史模型时更换gpu卡号,map_location设置

pytorch 在保存训练后模型的时候,会把训练过程中使用的设备号(例如gpu卡号cuda:0cpu)也一并保存下来。当pytorch重新载入历史模型时,模型默认根据训练时的设备卡号,把权值载入到相应的卡号上。

然而,有的时候测试过程和训练过程的设备情况是不一致的。
举个例子,A主机有四块GPU卡,然后我们用cuda:3 训练模型,并保存模型。
在测试时候,我们需要在客户的B主机跑模型,但是B主机只有一块gpu卡:cuda:0

如果按照默认方式载入模型的话,pytorch会报找不到gpu设备,或其他一些错误。

此时,载入的时候需要做一个变换,为torch.load指定gpu设备的映射方式:

根据pytorch的文档,在加载模型的时候,可以指定将模型的tensor加载到特定目标GPU上。
加载方法有:

>>> torch.load('tensors.pt')
# 1. Load all tensors onto the GPU 0
>>> torch.load('tensors.pt', map_location=torch.device('cuda:0'))
# 2. Load all tensors onto GPU 1
>>> torch.load('tensors.pt', map_location=lambda storage, loc: storage.cuda(1))
# 3. Map tensors from GPU 1 to GPU 0
>>> torch.load('tensors.pt', map_location={'cuda:1':'cuda:0'})

实测发现:
方法1 根本就没有加载到目标卡,模型原来在什么卡训练的,还是加载到原来的旧卡号上,因此指定失败。
方法3,代码之间就报错,location.startswith(‘cuda’): AttributeError: ‘NoneType’ object has no attribute ‘startswitch’,分析代码发现这是torch自己的bug! 坑爹的。
方法2: 可以正常的把tensor都加载到cuda1上。

  • 6
    点赞
  • 10
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值