PyTorch学习笔记（七）——两种方法保存训练好的神经网络

将代码和解释放在了一起：

# coding=gbk
import torch
import matplotlib.pyplot as plt

# 训练好了一个模型, 保存它, 留到下次要用的时候直接提取直接用
# 数据，以回归的模型作为例子
x = torch.unsqueeze(torch.linspace(-1, 1, 100), dim=1)  # x data (tensor), shape=(100, 1)
y = x.pow(2) + 0.2*torch.rand(x.size())  # noisy y data (tensor), shape=(100, 1)

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(500):
        prediction = net1(x)
        loss = loss_func(prediction, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    # 画出网络一结果图
    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)

    # 保存网络
    torch.save(net1, 'net.pkl')  # 保存整个网络，包括网络额结构、参数，用.pkl格式
    torch.save(net1.state_dict(), 'net_params.pkl') # 只保存网络的参数，用.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():
    # 先构建“网络三”的结构，然后再将之前保存的“net_params.pkl”填入
    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()
'''
总结：用Pytorch有两种保存网络的方法，第一种是保存整个网络，第二种是只保存参数
。第一种保存的文件相对较大，但是提取使用是较方便；第二种保存的文件相对较小，但是
提取使用的时候要先定义网络结构。
'''

 结果可视化，三个模型几乎是一样的，说明了两种保存模型的方法效果是一致的。