pytorch系列(四)：模型的保存和提取

ゞ浴缸里的玫瑰 2023-07-10 08:58 80阅读 0赞

import torch
    import torch.nn.functional as f
    
    #构造假数据
    # x=torch.squeeze(torch.linspace(-1,1,100))#shape为[100]的数组
    x=torch.unsqueeze(torch.linspace(0-1,1,100),dim=1)#shape为(100,1)的tensor
    y=x.pow(2)+0.2*torch.rand(x.size())
    
    #快速搭建网络
    def build_Net():
        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_fun= torch.nn.MSELoss()
        
        
        #训练
        for t in range(100):
            prediction=net1(x)#网络输出值  预测值
            loss=loss_fun(prediction,y)#计算损失值
            optimizer.zero_grad()#梯度清零
            loss.backward()#反向传播
            optimizer.step()#所有参数更新
            
        return net1
    
    net1=build_Net()
    
    #保存模型的两种方式
    #`1、
    torch.save(net1,"net.pkl")#保存整个网络
    #第二种
    torch.save(net1.state_dict(),"net_params.pkl")#只保存模型的参数
    
    
    #提取网络
    #第一种方法，提取整个网络
    def restore_Net():
        #把整个网络提取到net2中
        net2=torch.load("net.pkl")
        prediction=net2(x)
        return prediction
    
    prediction1=restore_Net()
    
    #第二种方式  只提取网络参数  然后直接放到新建的模型中去即可
    #新建一个无参数模型
    def new_Net():
        net3=torch.nn.Sequential(
            torch.nn.Linear(1,10),
            torch.nn.ReLU(),
            torch.nn.Linear(10,1)
        )
        #net3加载网络参数
        net3.load_state_dict(torch.load("net_params.pkl"))
        prediction=net3(x)
        return prediction
        
    prediction2=new_Net()
    
    #作图
    import matplotlib.pyplot as plt
    
    
    #原始数据和两种方法得到的预测数据   结果一致
    plt.scatter(x.data.numpy(),y.data.numpy(),c="red")
    plt.plot(x.data.numpy(),prediction1.data.numpy(),'bo')
    plt.plot(x.data.numpy(),prediction2.data.numpy(),'g+')