使用 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

可以发现，和之前的结果一样，再看最终绘制的图片：

图片也一致，说明模型加载成功