[深度學習] Pytorch（三）—— 多/單GPU、CPU，訓練保存、加載模型參數問題

[深度學習] Pytorch（三）—— 多/單GPU、CPU，訓練保存、加載預測模型問題

上一篇實踐學習中，遇到了在多/單個GPU、GPU與CPU的不同環境下訓練保存、加載使用使用模型的問題，如果保存、加載的上述三類環境不同，加載時會出錯。就去研究了一下，做了實驗，得出以下結論：

多/單GPU訓練保存模型參數、CPU加載使用模型

#保存
PATH = 'cifar_net.pth'
torch.save(net.module.state_dict(), PATH)

#加載
net = Net()
net.load_state_dict(torch.load(PATH))

多GPU訓練模型、單GPU加載使用模型

#保存
PATH = 'cifar_net.pth'
torch.save(net.state_dict(), PATH)

#加載
net = Net()
net = nn.DataParallel(net)  #保存多GPU的，在加載時需要把網絡也轉成DataParallel的
net.to(device)  #放到GPU上
net.load_state_dict(torch.load(PATH))

# 然后測試數據也需要放到GPU上
images, labels = images.to(device), labels.to(device)

多GPU訓練保存模型參數、多GPU加載使用模型

#保存
PATH = 'cifar_net.pth'
torch.save(net.state_dict(), PATH)

#加載
net = Net()
net = nn.DataParallel(net)  #保存多GPU的，在加載時需要把網絡也轉成DataParallel的
net.to(device)  #放到GPU上
net.load_state_dict(torch.load(PATH))

# 然后測試數據也需要放到GPU上
images, labels = images.to(device), labels.to(device)

可以看到，單GPU和多GPU加載數據的方法其實是一樣的，經運行驗證，只要按上述代碼寫，有多個GPU就調用多個，只有一個就調用一個。

另外，保存、加載網絡模型有三種不同的做法

1.保存整個網絡模型
2.只保存模型參數（我們用的就是這種）
3.自定義保存

詳細方法，請參考：https://blog.csdn.net/Code_Mart/article/details/88254444