Pytorch--训练并保存的模型利用 torch.load() 在GPU、CPU上加载

 

# 直接加载模型
model.load_state_dict(torch.load('./data/my_model.pkl'))

#GPU训练的模型加载到CPU上：
model.load_state_dict(torch.load('./data/my_model.pkl', map_location=lambda storage, loc: storage))

#加载到GPU1上：
model.load_state_dict(torch.load('./data/my_model.pkl', map_location=lambda storage, loc: storage.cuda(1)))

#从GPU1 移动到 GPU0：
model.load_state_dict(torch.load('./data/my_model.pkl', map_location={'cuda:1':'cuda:0'}))