pytorch在CPU和GPU上加载模型

 

 

pytorch允许把在GPU上训练的模型加载到CPU上，也允许把在CPU上训练的模型加载到GPU上。


CPU->CPU,GPU->GPU

 

 

torch.load('gen_500000.pkl')

 

 

GPU->CPU

 

 

torch.load('gen_500000.pkl', map_location=lambda storage, loc: storage)

或：

torch.load('xxx.pkl', map_location='cpu')  # pytorch0.4.0及以上版本

 

CPU->GPU1

 

 

torch.load('gen_500000.pkl', map_location=lambda storage, loc: storage.cuda(1))