Pytorch学习笔记【7】：快速搭建神经网络，神经网络的保存与加载

 还是老话，注意看代码 注释，解析都在注释里

一. 快速搭建神经网络

1. 代码

import torch
import torch.nn.functional as F

"""
此段代码主要诠释了如何快速搭建神经网络以及两种搭建方法的输出比较，
我们可以看到输出略有不同，但是两种方式搭建的神经网络的效果是一模一样的。
"""

class Net(torch.nn.Module):
    def __init__(self,n_feature,n_hidden,n_output):
        super(Net, self).__init__()
        self.hidden = torch.nn.Linear(n_feature,n_hidden)
        self.predict = torch.nn.Linear(n_hidden,n_output)

        def forward(self, x):
            x = F.relu(self.hidden(x)) # 隐藏层的输出要经过激活函数
            x = self.predict(x)
            return x

net1 = Net(1,10,1)

# 另一种快速搭建神经网络的方法
net2 = torch.nn.Sequential(
    torch.nn.Linear(1,10),
    torch.nn.ReLU(),
    torch.nn.Linear(10,1)
)

print(net1)
print(net2)

2.  运行结果

二.  神经网络的保存和加载

1. 代码

import torch
from torch.autograd import Variable
import matplotlib.pyplot as plt

"""
这段代码主要介绍如何保存我们训练好的神经网络以及如何加载我们保存好的神经网络
"""

# 造一些数据
x = torch.unsqueeze(torch.linspace(-1,1,100),dim=1)
y = x.pow(2) + 0.2*torch.rand(x.size())
x,y = Variable(x,requires_grad=False),Variable(y,requires_grad=False)

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

    # 训练
    for t in range(100):
        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') # 保存整个网络
    torch.save(net1.state_dict(),'net_params.pkl') # 只保存他的参数，一般这种方式后面加载会快一点，用的比较多


def load_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 load_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(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()
load_net()
load_params()

 2. 运行结果