训练好了一个模型, 我们当然想要保存它, 留到下次要用的时候直接提取直接用, 这就是这节的内容啦. 我们用回归的神经网络举例实现保存提取.
我们快速地建造数据, 搭建网络:
import torch
from torch.autograd import Variable
import matplotlib.pyplot as plt
#假数据
# unsqueeze对数据维度进行扩充
x = torch.unsqueeze(torch.linspace(-1,1,100),dim=1)
# print(x.size())
y = x.pow(2) + 0.2*torch.rand(x.size())
# print(y.shape)
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(100):
predict = net1(x)
loss = loss_func(predict,y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
plt.figure(1,figsize=(10,3))
plt.subplot(121)
plt.title("net1")
plt.scatter(x.data.numpy(),y.data.numpy())
plt.plot(x.data.numpy(),predict.data.numpy(),'r-')
# 保存整个网络
torch.save(net1,'net.pkl')
torch.save(net1, ‘net.pkl’) 这一行代码就是对整个网络模型的保存,这种方式将会提取整个神经网络, 网络大的时候可能会比较慢.
def restore_net():
net2 = torch.load('net.pkl')
prediction = net2(x)
plt.subplot(122)
plt.title("net2")
plt.scatter(x.data.numpy(),y.data.numpy())
plt.plot(x.data.numpy(),prediction.data.numpy(),'r-')
plt.show()
net2 = torch.load(‘net.pkl’) 就是对训练网络的提取
调用上面建立的几个功能, 然后出图.
#保存net1(整个网络)
save()
#提取到整个网络
restore_net()
接下来,就是保存网络参数,然后根据网络重新构建网络
import torch
from torch.autograd import Variable
import matplotlib.pyplot as plt
#假数据
# unsqueeze对数据维度进行扩充
x = torch.unsqueeze(torch.linspace(-1,1,100),dim=1)
# print(x.size())
y = x.pow(2) + 0.2*torch.rand(x.size())
# print(y.shape)
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(100):
predict = net1(x)
loss = loss_func(predict,y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
plt.figure(1,figsize=(10,3))
plt.subplot(121)
plt.title("net1")
plt.scatter(x.data.numpy(),y.data.numpy())
plt.plot(x.data.numpy(),predict.data.numpy(),'r-')
# 只保留网络中的参数(速度快,占内存少)
torch.save(net1.state_dict(),'net_params.pkl')
torch.save(net1.state_dict(),‘net_params.pkl’)这一行代码就是只保留网络中的参数(速度快,占内存少)
这种方式将会提取所有的参数, 然后再放到你的新建网络中.
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(122)
plt.title("net3")
plt.scatter(x.data.numpy(),y.data.numpy())
plt.plot(x.data.numpy(),prediction.data.numpy(),'r-')
这里就需要重新构建之前的网络,然后把保存的参数放进新的网络中
#保存net1(只有参数)
save()
#提取网络参数, 复制到新网络
restore_params()