代码如下:
#实现网络的保存和提取 import torch from torch.autograd import Variable import matplotlib.pyplot as plt #设置随机种子实现结果复现,在神经网络中,参数默认是进行随机初始化的。 # 不同的初始化参数往往会导致不同的结果,当得到比较好的结果时我们通常希望这个结果是可以复现的, # 在pytorch中,通过设置随机数种子也可以达到这么目的 torch.manual_seed(1) #生成数据 x = torch.unsqueeze(torch.linspace(-1, 1, 100), dim=1) y = x.pow(2) + 0.2*torch.rand(x.size()) x,y = Variable(x, requries_grad=False), Variable(y,requries_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(net.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() torch.save(net1, 'net1.pkl') #保存整个神经网络到net1.pkl中 torch.save(net1.state_dict(), 'net1_paras.pkl') #保存网络里的参数到net1_paras.pkl中 #画图 #提取方式1 #提取整个网络 def restore_net(): net2 = torch.load('net1.pkl') #提取方式2 #先建一个一样的网络,再把保存的参数放进去 def restore_paras(): net3 = torch.nn.Sequential( #建立和net1一样的层,不过参数肯定不同 torch.nn.Linear(1, 10), torch.nn.ReLU(), torch.nn.Linear(10, 1) ) net3.load_state_dict(torch.load('net1_paras.pkl'))
本站文章如无特殊说明,均为本站原创,如若转载,请注明出处:pytorch实现网络的保存和提取 - Python技术站