pytorch神经网络 训练好的模型保存提取

要点

训练好了一个模型, 我们当然想要保存它, 留到下次要用的时候直接提取直接用, 这就是这节的内容啦. 我们用回归的神经网络举例实现保存提取.

对整个网络保存

我们快速地建造数据, 搭建网络:

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()

可以看出net1和net2的图一模一样

接下来，就是保存网络参数，然后根据网络重新构建网络

只保留网络中的参数

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()

可以看出，net1和net3一模一样