pytorch之 sava_reload_model

 import torch
 import matplotlib.pyplot as plt
 
 # torch.manual_seed(1)    # reproducible
 
 # 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)
 
 # 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(
         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()
 
     # 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')
     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():
     # 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)
     )
 
     # copy net1's parameters into 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()
 
 # save net1
 save()
 
 # restore entire net (may slow)
 restore_net()
 
 # restore only the net parameters
 restore_params()