pytorch深度学习：模型的保存和提取

在学习完搭建神经网络模型后，有时需要对模型进行保存和提取，这节课就是讲述如何在pytorch中实现。

1. 问题的提出

我们使用回归这节课建立的神经网络来作为要保存的模型，具体的模型建立过程可以参考我这篇博客。它是对一元二次方程进行拟合的模型 (net1)。

2. 模型的保存

神经网络模型的保存有两种方法：

• 保存整个网络模型。
• 只保存网络模型的参数（速度快，占用内存少）。

# 1.保存整个模型
torch.save(net1,'net.pkl')
# 2.只保存模型参数
torch.save(net1.state_dict(),'net_parms.pkl')

3. 模型的提取

针对两种保存方法，会有两种模型提取的方法。

• 提取整个网络模型 （net2）。
• 只提取模型参数（需要先建立一个net3和net1结构一致）。

# 1.提取整个模型
net2 = torch.load('net.pkl')

# 2.只提取模型参数,要首先建立一个相同格式的网络
net3 = Net(1,10,1)
net3.load_state_dict(torch.load('net_parms.pkl'))

4. 显示结果

将net1、net2、net3的模型分别作图，可以发现它们的结果完全一致。

代码参考了莫烦大神的教程，完整代码放在这里：

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt

# 建立数据:拟合y = x^2曲线
x = torch.unsqueeze(torch.linspace(-1, 1, 100), dim=1)  # x data (tensor), shape=(100, 1)
y = x.pow(2) + 0.2*torch.rand(x.size())                 # noisy y data (tensor), shape=(100, 1)

# 建立模型
class Net(nn.Module):
    def __init__(self,n_feature,n_hidden,n_output):
        super(Net,self).__init__()
        self.hidden = nn.Linear(n_feature,n_hidden)
        self.predict = nn.Linear(n_hidden,n_output)
        
    def forward(self,x):
        x = F.relu(self.hidden(x))
        x = self.predict(x)
        
        return x

net1 = Net(1,10,1)
optimizer = torch.optim.SGD(net1.parameters(),lr=0.1)
loss_func = nn.MSELoss()

# 训练模型
for i in range(100):
    prediction = net1(x)
    loss = loss_func(prediction,y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
# 保存模型
# 1.保存整个模型
torch.save(net1,'net.pkl')
# 2.只保存模型参数
torch.save(net1.state_dict(),'net_parms.pkl')

# 提取模型
# 1.提取整个模型
net2 = torch.load('net.pkl')
# 2.只提取模型参数,要首先建立一个相同格式的网络
net3 = Net(1,10,1)
net3.load_state_dict(torch.load('net_parms.pkl'))

# 画图
plt.figure(1, figsize=(10, 3))
plt.subplot(131)
plt.title('net1')
plt.scatter(x.data.numpy(), y.data.numpy())
plt.plot(x.data.numpy(), prediction.data.numpy(), 'r-', lw=5)

plt.subplot(132)
plt.title('net2')
plt.scatter(x.data.numpy(), y.data.numpy())
plt.plot(x.data.numpy(), net2(x).data.numpy(), 'r-', lw=5)

plt.subplot(133)
plt.title('net3')
plt.scatter(x.data.numpy(), y.data.numpy())
plt.plot(x.data.numpy(), net3(x).data.numpy(), 'r-', lw=5)