pytroch--保存和提取神经网络

最新推荐文章于 2024-06-23 16:19:59 发布

原创最新推荐文章于 2024-06-23 16:19:59 发布 · 334 阅读

0 ·

CC 4.0 BY-SA版权

文章标签：

#神经网络 #保存提取

本文介绍如何在PyTorch中保存已训练的神经网络模型，并详细讲解了模型的加载过程，确保可以恢复训练状态。

我们想要保存神经网络的模型，就需要知道保存和提取我们已经训练好的模型：

在这里插入代码片
import torch
import numpy as np
import matplotlib.pyplot as plt
from torch.nn.functional as F  #获得神经网络的激励函数
from torch.autograd import Variable  #变量（将张量数据变量化，神经网络只能处理变量）

x = torch.unsqueeze(torch.linspace(-1,1,100),dim=1)  #使用unsqueeze将数据转换为二维数据，需要用上dim参数
y = x.pow(2) + 0.2*torch.rand(x.size())
x, y = Variable(x,requires_grad=False), Variable(y,requires_grad=False)


#搭神经网络
def save():
	net = 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 = net(x)
		loss = loss_func(prediction, y)
		optimizer.zero_grad()
		loss.backward()
		optimizer.step()
	
		torch.save(net,'net.pkl')  #保存在整个神经网络
		torch.save(net.state_dict(),'net_params.pkl')    #只保存神经网络的参数
		
		#以下为画图代码：
		plt.figure(figsize=(10,3))
		plt.sunplot(131)
		plt.title('net')
	    plt.scatter(x.data.numpy(), y.data.numpy())
	    plt.plot(x.data.numpy(), prediction.data.numpy(), 'r-', lw=5)
	    plt.show()


#提取整个神经网络
def restore_net():
	net1 = torch.load('net.pkl')
	prediction = net1(x)
	plt.sunplot(132)
	plt.title('net1')
	plt.scatter(x.data.numpy(), y.data.numpy())
	plt.plot(x.data.numpy(), prediction.data.numpy(), 'r-', lw=5)
	plt.show()

#提取神经网络的参数（这个方法比较快，但是需要重新建立一个一样的神经网络）
def restore_params():
	net2 = torch.nn.Sequential(
		torch.nn.Linear(1,10),
		torch.nn.ReLU(),
		torch.nn.Linear(10,1)
	)
	net2 = torch.load_state_dict(torch.load('net_params.pkl'))
	prediction = net2(x)
	plt.sunplot(133)
	plt.title('net2')
	plt.scatter(x.data.numpy(), y.data.numpy())
	plt.plot(x.data.numpy(), prediction.data.numpy(), 'r-', lw=5)
	plt.show()