[pytorch学习笔记四] 保存和提取训练好的网络(含实例详解)

学自莫凡PYTHON

1.方法

保存：

保存训练好的整个网络：torch.save(net1,'net1.pth')
只保存网络的参数：torch.save(net1.state_dict(),'net1_params.pth')

提取：

调用整个训练好的网络：net2 = torch.load('net1.pth')
只调用网络的参数：net3.load_state_dict(torch.load('net1_params.pth'))

其中，net1是训练好的网络的名称，’ '内部是保存的文件名称（后缀是.pth或者.pkl）

注意：
（1）只保存网络中的参数速度快, 占内存少，推荐使用。但是调用网络的参数时，新网络需要自己定义网络，再使用上面的调用指令。其中的参数名称与结构要与保存模型中的一致，可以是部分网络比如只使用VGG的前几层，相对灵活，便于对网络进行修改。
（2）调用整个网络则无需自定义网络可以直接调用。保存时已把网络结构保存，比较死板，不能调整网络结构。

2.实例

2.1 实验结果

以我的[pytorch学习笔记二]数据的拟合为例，将训练好的net1网络保存好，net2调用的整个网络，net3只调用了net1的参数，最终拟合的效果一模一样。

2.2完整代码

# 1.导入必要的模块
import torch
import torch.nn.functional as F   # F中包含很多函数比如激励函数
import matplotlib.pyplot as plt #用于绘图

# 2.生成要拟合的数据点
x = torch.unsqueeze(torch.linspace(-1,1,100),dim=1)
y = x.pow(2) + 0.2 * torch.rand(x.size())

# 3.搭建、训练、保存网络net1
def save():
    net1 = torch.nn.Sequential(
        torch.nn.Linear(1, 8),
        torch.nn.ReLU(),
        torch.nn.Linear(8, 1)
    )

    optimizer = torch.optim.SGD(net1.parameters(), lr=0.5)
    loss_func = torch.nn.MSELoss()

    for t in range(100):
        prediction = net1(x)
        loss = loss_func(prediction, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    # 保存训练好的整个网络
    torch.save(net1, 'net1.pth')
    # 只保存训练好的网络的参数
    torch.save(net1.state_dict(), 'net1_params.pth')

    # 绘图，可视化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)

# 调用整个网络（包括计算图、参数等等）
def restore_net():
    net2 = torch.load('net1.pth')
    # 调用网络时，要设置输入
    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():
    # 只调用网络的参数时，需要提前搭建和net1网络相同的架构
    net3 = torch.nn.Sequential(
        torch.nn.Linear(1, 8),
        torch.nn.ReLU(),
        torch.nn.Linear(8, 1)
    )
    net3.load_state_dict(torch.load('net1_params.pth'))
    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()