保存方法
先建立一个网络 net1
torch.save(net1, 'net.pkl') #保存全部网络 torch.save(net1.state_dict(), 'net_params.pkl') #保存参数
提取方法1 提取全部网络
def restore_net(): # restore entire net1 to net2 net2 = torch.load('net.pkl') #直接提取net1 prediction2 = net2(x)
提取方法2 提取参数
需要自己再建立一个跟要提取的网络 一模一样的网络(本文是net3)
然后提取参数
net3 = torch.nn.Sequential( torch.nn.Linear(1, 10), torch.nn.ReLU(), torch.nn.Linear(10, 1) ) # copy net1'parameters into net3 net3.load_state_dict(torch.load('net_params.pkl')) prediction3 = net3(x)
全部代码如下
import torch from torch.autograd import Variable import matplotlib.pyplot as plt # torch.manual_seed(1) # 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, 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') torch.save(net1.state_dict(), 'net_params.pkl') 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'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() # swve net1 save() # restore entire net(may slow) restore_net() # restore only the net parameters restore_params()