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