pytorch模型的保存与加载

我们先创建一个模型，使用的是pytorch笔记——简易回归问题_刘文巾的博客-CSDN博客 的主体框架，唯一不同的是，我这里用的是torch.nn.Sequential来定义模型框架，而不是那篇博客里面的类。

1 保存与加载之前的部分

#导入库
import torch

#数据集
x=torch.linspace(-1,1,100).reshape(-1,1)
y=x*x+0.2*torch.rand(x.shape)

#定义模型（Sequential比类简明了很多）
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.2)

loss_func=torch.nn.MSELoss()

#进行训练
for epoch in range(1000):
    prediction=net(x)
    
    loss=loss_func(prediction,y)
    
    optimizer.zero_grad()#清空上一轮参数优化的参与梯度
    
    loss.backward()#损失函数反向传播
    
    optimizer.step()#梯度更新

#打印模型里面的参数
for a,b in enumerate(net.parameters()):
    print('no:',a,'\n',b)
'''
一共有四个参数，分别对应的是每一层的w和b
no: 0 
 Parameter containing:
tensor([[ 1.0588],
        [ 0.1654],
        [ 0.8578],
        [-0.8756],
        [-1.0935],
        [ 0.7588],
        [ 0.9043],
        [-0.0723],
        [-0.4335],
        [ 0.3010]], requires_grad=True)
no: 1 
 Parameter containing:
tensor([-0.5888,  0.9550, -0.1572, -0.2610, -0.4367,  0.3084, -0.3802, -0.3834,
         0.6192,  0.3012], requires_grad=True)
no: 2 
 Parameter containing:
tensor([[ 0.3397,  0.0161,  0.7238,  0.6869,  0.5263,  0.1717,  0.6978,  0.1012,
          0.3311, -0.2264]], requires_grad=True)
no: 3 
 Parameter containing:
tensor([-0.0836], requires_grad=True)
'''

2 存储与加载（方法1）——直接保存模型

我们使用torch.save直接保存模型

torch.save(net,'net.pkl')

加载模型的时候，直接torch.load即可（可以看到net2参数和net是一样的）

net2=torch.load('net.pkl')
for a,b in enumerate(net2.parameters()):
    print('no:',a,'\n',b)

'''
no: 0 
 Parameter containing:
tensor([[ 1.0588],
        [ 0.1654],
        [ 0.8578],
        [-0.8756],
        [-1.0935],
        [ 0.7588],
        [ 0.9043],
        [-0.0723],
        [-0.4335],
        [ 0.3010]], requires_grad=True)
no: 1 
 Parameter containing:
tensor([-0.5888,  0.9550, -0.1572, -0.2610, -0.4367,  0.3084, -0.3802, -0.3834,
         0.6192,  0.3012], requires_grad=True)
no: 2 
 Parameter containing:
tensor([[ 0.3397,  0.0161,  0.7238,  0.6869,  0.5263,  0.1717,  0.6978,  0.1012,
          0.3311, -0.2264]], requires_grad=True)
no: 3 
 Parameter containing:
tensor([-0.0836], requires_grad=True)
'''

3 存储和加载（方法2）——保存模型参数

保存的话存模型的参数

torch.save(net.state_dict(),'net_params.pkl')

加载的话，我们得先重新声明一个新的神经网络结构（用Sequential和用类都可以，有了这个新的神经网络后，才可以把参数传进去）【因为在声明新的神经网络之前，我们现在存的内容即使加载出来了，也不知道这些参数对应的结构是什么】

#声明一个新的net
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') )

for a,b in enumerate(net3.parameters()):
    print('no:',a,'\n',b)
'''
和net也是一样的
no: 0 
 Parameter containing:
tensor([[ 1.0588],
        [ 0.1654],
        [ 0.8578],
        [-0.8756],
        [-1.0935],
        [ 0.7588],
        [ 0.9043],
        [-0.0723],
        [-0.4335],
        [ 0.3010]], requires_grad=True)
no: 1 
 Parameter containing:
tensor([-0.5888,  0.9550, -0.1572, -0.2610, -0.4367,  0.3084, -0.3802, -0.3834,
         0.6192,  0.3012], requires_grad=True)
no: 2 
 Parameter containing:
tensor([[ 0.3397,  0.0161,  0.7238,  0.6869,  0.5263,  0.1717,  0.6978,  0.1012,
          0.3311, -0.2264]], requires_grad=True)
no: 3 
 Parameter containing:
tensor([-0.0836], requires_grad=True)
'''

4 两种方法的比较

存参数的文件占用的空间少一点，这个在目前这种比较简单的模型可能还看不出来。对于那种大的模型，省下来的空间还是蛮多的。