将代码和解释放在了一起:
# coding=gbk
import torch
import matplotlib.pyplot as plt
# 训练好了一个模型, 保存它, 留到下次要用的时候直接提取直接用
# 数据,以回归的模型作为例子
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)
def save():
# 用快速建立网络的方法建立“网络一”
net1 = torch.nn.Sequential(
torch.nn.Linear(1, 10),
torch.nn.ReLU(),
torch.nn.Linear(10, 1)
)
optimizer = torch.optim.SGD(net1.parameters(), lr=0.5)#优化器
loss_func = torch.nn.MSELoss()#损失函数
for t in range(500):
prediction = net1(x)
loss = loss_func(prediction, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
# 画出网络一结果图
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)
# 保存网络
torch.save(net1, 'net.pkl') # 保存整个网络,包括网络额结构、参数,用.pkl格式
torch.save(net1.state_dict(), 'net_params.pkl') # 只保存网络的参数,用.pkl格式
# 只保存参数的形式虽然保存时更简洁,但是调用网络时要重新搭建网络
def restore_net():
# 将保存的整个网络提取出来,作为“网络二”
net2 = torch.load('net.pkl')
prediction = net2(x)
# 画出网络二结果图
plt.subplot(132)
plt.title('Net2')
plt.scatter(x.data.numpy(), y.data.numpy())
plt.plot(x.data.numpy(), prediction.data.numpy(), 'r-', lw=5)
def restore_params():
# 先构建“网络三”的结构,然后再将之前保存的“net_params.pkl”填入
net3 = torch.nn.Sequential(
torch.nn.Linear(1, 10),
torch.nn.ReLU(),
torch.nn.Linear(10, 1)
)
net3.load_state_dict(torch.load('net_params.pkl'))
prediction = net3(x)
#画出网络三结果图
plt.subplot(133)
plt.title('Net3')
plt.scatter(x.data.numpy(), y.data.numpy())
plt.plot(x.data.numpy(), prediction.data.numpy(), 'r-', lw=5)
plt.show()
# 保存网络一
save()
# 保存网络二
restore_net()
# 保存网络三
restore_params()
'''
总结:用Pytorch有两种保存网络的方法,第一种是保存整个网络,第二种是只保存参数
。第一种保存的文件相对较大,但是提取使用是较方便;第二种保存的文件相对较小,但是
提取使用的时候要先定义网络结构。
'''
结果可视化,三个模型几乎是一样的,说明了两种保存模型的方法效果是一致的。