pytorch 多卡並行計算保存模型和加載模型 （遺漏module的解決）

今天使用了多卡進行訓練，保存的時候直接是用了下面的代碼：

torch.save(net.cpu().state_dict(),'epoch1.pth')

我在測試的時候，想要加載這個訓練好的模型，但是報錯了，說是字典中的關鍵字不匹配，我就將新創建的模型，和加載的模型中的關鍵字都打印了出來，發現夾雜的模型的每個關鍵字都多了module. 。解決方式爲：

pre_dict = torch.load('./epoch1.pth')
new_pre = {}
for k,v in pre_dict.items():
    name = k[7:]
    new_pre[name] = v

net.load_state_dict(new_pre)

這就相當於是把不同的關鍵字都設置成相同的關鍵字，也將參數加載了進來。