本文共 1769 字，大约阅读时间需要 5 分钟。

学习莫烦pytorch视频，训练好的网络进行保存和读取

#qd.py#classification.pyimport numpy as npimport torchimport torch.nn.functional as Ffrom torch.autograd import Variableimport matplotlib.pyplot as pltimport mathimport pdbtorch.manual_seed(1)#CPU设置种子用于生成随机数，随机值确定x = torch.unsqueeze(torch.linspace(-1, 1, 100), dim=1)##torch.linspace(-1, 1, 100) 【-1,1】等分成100份的等差数列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.2)    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,'net.pkl')#entire  net整个网络    torch.save(net1.state_dict(), 'net_params.pkl') #parameters 只有参数    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('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():    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()

转载地址：http://awksi.baihongyu.com/