前言

在使用Pytorch訓練模型的時候，經常會有在GPU上儲存模型然後再CPU上執行的需求，在實驗的過程中發現在多GPU上訓練的Pytorch模型是不能在CPU上直接執行的，幾次遇到了這種問題，這裡研究和記錄一下。

模型的儲存與載入

例如我們建立了一個模型：

model = MyVggNet()

如果使用多GPU訓練,我們需要使用這行程式碼：

model = nn.DataParallel(model).cuda()

執行這個程式碼之後，model就不在是我們原來的模型，而是相當於在我們原來的模型外面加了一層支援CPU執行的外殼，這時候真正的模型物件為：real_model = model.module

torch.save(model_object, 'model.pkl')
model = torch.load('model.pkl')

1. 僅儲存和載入模型引數(推薦使用)

torch.save(model_object.state_dict(), 'params.pkl')
model_object.load_state_dict(torch.load('params.pkl'))

模型儲存與載入對應方式

1. 第一種方式

儲存使用：

real_model = model.module
torch.save(real_model.state_dict(),os.path.join(args.save_path,"cos_mnist_"+str(epoch+1)+"_weight.pth"))

cpu上載入使用：

args.weight=checkpoint/cos_mnist_10_weight.pth
map_location = lambda
 
 storage, loc: storage
model.load_state_dict(torch.load(args.weight,map_location=map_location))

2. 第二種方式

儲存使用：

real_model = model.module
save_model(real_model, os.path.join(args.save_path,"cos_mnist_"+str(epoch+1)+"_weight_cpu.pth"))
# 自定義的函式
def save_model(model,filename):
    state = model.state_dict()
    for key in state: state[key] = state[key].clone().cpu()
    torch.save(state, filename)

cpu上載入使用：

args.weight=checkpoint/cos_mnist_10_weight_cpu.pth
model.load_state_dict(torch.load(args.weight))

3. 第三種方式

儲存使用：

real_model = model.module
torch.save(real_model, os.path.join(args.save_path,"cos_mnist_"+str(epoch+1)+"_whole.pth"))

cpu上載入使用：

args.weight=checkpoint/cos_mnist_10_whole.pth
map_location = lambda storage, loc: storage
model = torch.load(args.weight,map_location=map_location)

參考文獻