使用 pytorch 保存和加載訓練好的模型

本篇博客講述如何使用 pytorch 保存訓練好的神經網絡和如何將訓練好的神經網絡加載進來以便使用。

定義網絡結構

這裏使用最簡單的一個結構，兩個線性的全連接層，有激活函數，用來擬合二維空間上的一些點

"""
net.py  用於定義網絡的結構
"""
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.layer1=nn.Linear(1,3)
        self.layer2=nn.Linear(3,1)
    def forward(self,x):
        x=self.layer1(x)
        x=torch.relu(x)
        x=self.layer2(x)
        return x

訓練神經網絡並且保存

使用網絡擬合一個散點圖，經過訓練之後的模型保存成爲 pth 文件

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
from net import Net
# 訓練兩千次
num_epoches=2000
# 學習率定義爲0.01
learning_rate=0.01
# 創造一堆散點（用於訓練神經網絡）
x_train=np.array([
    3.3,4.4,5.5,6.71,6.93,4.168,
    9.779,6.182,7.59,2.167,7.042,
    10.791,5.313,7.997,3.1
],dtype=np.float32).reshape(-1,1)
y_train=np.array([
    1.7,2.76,2.09,3.19,1.694,1.573,
    3.366,2.596,2.53,1.221,2.827,
    3.465,1.65,2.904,1.3
],dtype=np.float32).reshape(-1,1)
# 創建一個模型
model=Net()
# 使用平方差均值來作爲損失函數
loss_fun=nn.MSELoss()
optimizer=torch.optim.SGD(model.parameters(),lr=learning_rate)

# 開始訓練：
inputs=torch.from_numpy(x_train)
targets=torch.from_numpy(y_train)
for epoch in range(num_epoches):
    output=model(inputs)
    loss=loss_fun(output,targets)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if (epoch+1) % 10==0:
        print("Epoch {} / {},loss {:.4f}".format(epoch+1,num_epoches,loss.item()))

# 保存訓練的模型：
torch.save(model.state_dict(),"Linear.pth")

# 打印最終的損失值
output=model(inputs)
loss=loss_fun(output,targets)
print(loss.item())

# 繪製經過網絡之後的預測圖形和原本的圖形，使用matplot 繪製曲線
predicted=model(torch.from_numpy(x_train)).detach_().numpy()
plt.plot(x_train,y_train,'ro',label="Origin data")
plt.plot(x_train,predicted,label="Fitted data")
plt.legend()
plt.show()

這段代碼的輸出爲：

0.16891564428806305
Process finished with exit code 0

得到的圖像爲：

經過上述代碼，我們就已經將 模型保存在了本地的 Linear.pth 文件中。接下來就是讀取該文件，並將訓練好的參數全部加載進去

讀取預訓練的模型並使用

使用和訓練模型一樣的數據，如果確實成功導入了，那麼最終的loss值和所繪製的圖形應該完全一樣，接下來是驗證的代碼：

import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import numpy as np
from net import Net
# 創建一個一模一樣的模型
model=Net()
# 加載預訓練模型的參數
model.load_state_dict(torch.load("Linear.pth"))

# 使用和訓練數據一樣的數據
x_train=np.array([
    3.3,4.4,5.5,6.71,6.93,4.168,
    9.779,6.182,7.59,2.167,7.042,
    10.791,5.313,7.997,3.1
],dtype=np.float32).reshape(-1,1)

y_train=np.array([
    1.7,2.76,2.09,3.19,1.694,1.573,
    3.366,2.596,2.53,1.221,2.827,
    3.465,1.65,2.904,1.3
],dtype=np.float32).reshape(-1,1)

# 計算loss 的值
inputs=torch.from_numpy(x_train)
target=torch.from_numpy(y_train)
output=model(inputs)
loss_fun=nn.MSELoss()
loss=loss_fun(output,target)
print(loss.item())

# 繪製圖形
plt.plot(x_train,y_train,'ro',label='origin data')
plt.plot(x_train,output.detach_().numpy(),label='Fitted data')
plt.legend()
plt.show()

上面程序的輸出結果是：

0.16891564428806305
Process finished with exit code 0

可以發現，和之前的結果一樣，再看最終繪製的圖片：

圖片也一致，說明模型加載成功