pytorch中 gpu与gpu、gpu与cpu 在load时相互转化操作
问题描述
有时在加载已训练好的模型时,会出现outofmemory的错误提示,但仔细检测使用的GPU卡并没有再用且内存也没有超出。
经查阅发现原来是训练模型时使用的GPU卡和加载时使用的GPU卡不一样导致的。个人感觉,因为pytorch的模型中是会记录有GPU信息的,所以有时使用不同的GPU加载时会报错。
解决方法
gpu之间的相互转换。即,将训练时的gpu卡转换为加载时的gpu卡。
torch.load('modelparameters.pth',map_location={'cuda:1':'cuda:0'})#gpu1-->gpu0
当然,你也可以直接将加载模型时使用的gpu卡改为和训练时的一样。但在多人使用一个服务器时,你想用的gpu卡已被使用,就需按上面方法转换gpu。
拓展
cpu-->cpu或gpu-->gpu
checkpoint=torch.load('modelparameters.pth')
model.load_state_dict(checkpoint)
cpu-->gpu0
torch.load('modelparameters.pth',map_location=lambdastorage,loc:storage.cuda(0))
gpu-->cpu
torch.load('modelparameters.pth',map_location=lambdastorage,loc:storage)
gpu1-->gpu0
torch.load('modelparameters.pth',map_location={'cuda:1':'cuda:0'})
补充知识:pytorchmodel.cuda()花费时间很长
如果pytorch在进行model.cuda()操作需要花费的时间很长,长到你怀疑GPU的速度了,那就是不正常的。
如果你用的pytorch版本是0.3.0,升级到0.3.1就好了!
以上这篇pytorch模型载入之gpu和cpu互转操作就是小编分享给大家的全部内容了,希望能给大家一个参考,也希望大家多多支持毛票票。