快速搭建网络与保存

import torch
import matplotlib.pyplot as plt

# torch.manual_seed(1)    # reproducible
# 用Pytorch有两种保存网络的方法，第一种是保存整个网络，第二种是只保存参数
# 。第一种保存的文件相对较大，但是提取使用是较方便；第二种保存的文件相对较小，但是
# 提取使用的时候要先定义网络结构。


# fake data
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)

# x = torch.unsqueeze(torch.linspace(1, 10, 10), dim=1)  # x data (tensor), shape=(100, 1)
# y = torch.unsqueeze(torch.linspace(10, 1, 10), dim=1)  # x data (tensor), shape=(100, 1)

# The code below is deprecated in Pytorch 0.4. Now, autograd directly supports tensors
# x, y = Variable(x, requires_grad=False), Variable(y, requires_grad=False)


def save():
    # save net1
    net1 = torch.nn.Sequential(# Sequential实现快速搭建 建立神经网络结构
        torch.nn.Linear(1, 10),
        torch.nn.ReLU(),
        torch.nn.Linear(10, 1)
    )
    print(net1)
    # 随机梯度下降优化 net.parameters()传入 net 的所有参数, 学习率为0.5
    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()

    # plot result
    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)

    # 2 ways to save the net两种保存方式
    torch.save(net1, 'net.pkl')  # save entire net
    torch.save(net1.state_dict(), 'net_params.pkl')   # save only the parameters # 只保存网络中的参数 (速度快, 占内存少)


def restore_net():
    # restore entire net1 to net2
    net2 = torch.load('net.pkl')#加载训练好的net.pkl网络为net2 直接用net2预测测试
    prediction = net2(x)

    # plot result
    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
    # restore only the parameters in net1 to net3
    net3 = torch.nn.Sequential(
        torch.nn.Linear(1, 10),
        torch.nn.ReLU(),
        torch.nn.Linear(10, 1)
    )
    # 将保存的参数复制到 net3
    # copy net1's parameters into net3
    #建立相同结构的神经网路 直接把已经保存的参数给新建的网络
    #后直接用net3进行预测测试
    net3.load_state_dict(torch.load('net_params.pkl'))
    prediction = net3(x)

    # plot result
    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()

 # 保存 net1 (1. 整个网络, 2. 只有参数)
# save net1
save()

# restore entire net (may slow)
restore_net()

# restore only the net parameters
restore_params()