训练好了一个模型,想要保存它,留到下次要用的时候直接提取直接用,这就是这次学习的主要内容。
一. 保存
1.建造数据
x = torch.unsqueeze(torch.linspace(-1, 1, 100), dim=1)
y = x.pow(2) + 0.2*torch.rand(x.size())
注:
- torch.unsqueeze(input, dim) :返回在指定位置插入的尺寸为1的新张量(与原张量共用相同数据)。
- torch.linspace(start,end,steps = 100…) :返回 介于start和end之间的等距一维张量。
2.搭建网络
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()
3.训练数据
for t in range(100):
prediction = net1(x)
loss = loss_func(prediction, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
4.保存
torch.save(net1, 'net.pkl') # save entire net
torch.save(net1.state_dict(), 'net_params.pkl') # save only the parameters
二.提取
1.提取整个网络
net2 = torch.load('net.pkl')
prediction = net2(x)
2.提取网络参数
# restore only the parameters in net1 to net3
net3 = torch.nn.Sequential(
torch.nn.Linear(1, 10),
torch.nn.ReLU(),
torch.nn.Linear(10, 1)
)
# copy net1's parameters into net3
net3.load_state_dict(torch.load('net_params.pkl'))
prediction = net3(x)
注:提取参数的方法,需要知道网络的形式,并需要先搭建一个相同的网络。
三.画图
1.net1
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)
2.net2
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)
3.net3
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()
四 .结果
注:
- 可以将上面的步骤定义为功能模块,避免代码的重复复制粘贴。
- 学习参考pytorch官网和莫烦python。
- 完整代码参考这里