pytorch中 gpu与gpu、gpu与cpu 在load时相互转化

问题描述

有时在加载已训练好的模型时，会出现  out of memory  的错误提示，但仔细检测使用的GPU卡并没有再用且内存也没有超出。

经查阅发现原来是训练模型时使用的GPU卡和加载时使用的GPU卡不一样导致的。个人感觉，因为pytorch的模型中是会记录有GPU信息的，所以有时使用不同的GPU加载时会报错。 

解决方法

gpu之间的相互转换。即，将训练时的gpu卡转换为加载时的gpu卡。

torch.load('modelparameters.pth', map_location={'cuda:1':'cuda:0'})  # gpu 1 --> gpu 0

当然，你也可以直接将加载模型时使用的gpu卡改为和训练时的一样。但在多人使用一个服务器时，你想用的gpu卡已被使用，就需按上面方法转换gpu。

拓展

cpu --> cpu 或 gpu --> gpu

checkpoint = torch.load('modelparameters.pth')

model.load_state_dict(checkpoint)

cpu --> gpu 0

torch.load('modelparameters.pth', map_location=lambda storage, loc: storage.cuda(0))

gpu --> cpu

torch.load('modelparameters.pth', map_location=lambda storage, loc: storage)

gpu 1 --> gpu 0

torch.load('modelparameters.pth', map_location={'cuda:1':'cuda:0'})