导入固定的模块和创建数据(和数据的回归问题相同)
import torch
import numpy as np
import matplotlib.pyplot as plt
from torch.autograd import Variable
import torch.nn.functional as F
torch.manual_seed(1)#同样的随机初始化种子
#fake data
x=torch.unsqueeze(torch.linspace(-1,1,100),dim=1)
y=x.pow(2)+0.2*torch.rand(x.size())
x,y=Variable(x,requires_grad=False),Variable(y,requires_grad=False)
save net
分为保存整个网络和保存网络参数
def save():
#save net1
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.02)
loss_func= torch.nn.MSELoss()
for t in range (1000):
prediction = net1(x)#F.softmax(out)将数值转换成概率
loss = loss_func(prediction,y)
optimizer.zero_grad()#把梯度降维0
loss.backward()#反向传递
optimizer.step()#用optimizer优化过程
torch.save(net1,'net.pkl')#保存整个网络
torch.save(net1.state_dict(),'net_params.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():
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()
最终结果
三种网络图像一致
附完整代码
主要注意点在提取参数net时,要用torch.nn.Sequential再进行一遍搭建
import torch
import numpy as np
import matplotlib.pyplot as plt
from torch.autograd import Variable
import torch.nn.functional as F
torch.manual_seed(1)#同样的随机初始化种子
#fake data
x=torch.unsqueeze(torch.linspace(-1,1,100),dim=1)
y=x.pow(2)+0.2*torch.rand(x.size())
x,y=Variable(x,requires_grad=False),Variable(y,requires_grad=False)
def save():
#save net1
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.02)
loss_func= torch.nn.MSELoss()
for t in range (1000):
prediction = net1(x)#F.softmax(out)将数值转换成概率
loss = loss_func(prediction,y)
optimizer.zero_grad()#把梯度降维0
loss.backward()#反向传递
optimizer.step()#用optimizer优化过程
torch.save(net1,'net.pkl')#保存整个网络
torch.save(net1.state_dict(),'net_params.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)
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():
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()