• 企业400电话
  • 微网小程序
  • AI电话机器人
  • 电商代运营
  • 全 部 栏 目

    企业400电话 网络优化推广 AI电话机器人 呼叫中心 网站建设 商标✡知产 微网小程序 电商运营 彩铃•短信 增值拓展业务
    pytorch模型的保存和加载、checkpoint操作

    其实之前笔者写代码的时候用到模型的保存和加载,需要用的时候就去度娘搜一下大致代码,现在有时间就来整理下整个pytorch模型的保存和加载,开始学习把~

    pytorch的模型和参数是分开的,可以分别保存或加载模型和参数。所以pytorch的保存和加载对应存在两种方式:

    1. 直接保存加载模型

    (1)保存和加载整个模型

    # 保存模型
    torch.save(model, 'model.pth\pkl\pt')    #一般形式torch.save(net, PATH) 
    # 加载模型
    model = torch.load('model.pth\pkl\pt')   #一般形式为model_dict=torch.load(PATH)

    (2)仅保存和加载模型参数(推荐使用,需要提前手动构建模型)

    速度快,占空间少

    # 保存模型参数
    torch.save(model.state_dict(), 'model.pth\pkl\pt')  #一般形式为torch.save(net.state_dict(),PATH)
    # 加载模型参数
    model.load_state_dict(torch.load('model.pth\pkl\pt') #一般形式为model_dict=model.load_state_dict(torch.load(PATH))

    state_dict() 是一个Python字典,将每一层映射成它的参数张量。注意只有带有可学习参数的层(卷积层、全连接层等),以及注册的缓存(batchnorm的运行平均值)在state_dict 中才有记录。state_dict同样包含优化器对象,存储了优化器的状态,所使用到的超参数。

    然而,在实验中往往需要保存更多的信息,比如优化器的参数,那么可以采取下面的方法保存:

    torch.save({'epoch': epochID + 1, 'state_dict': model.state_dict(), 'best_loss': lossMIN,
            'optimizer': optimizer.state_dict(),'alpha': loss.alpha, 'gamma': loss.gamma},
            checkpoint_path + '/m-' + launchTimestamp + '-' + str("%.4f" % lossMIN) + '.pth.tar')

    如下一个完整的使用model.state_dict()和optimizer.state_dict()例子:

    # 定义模型
    class TheModelClass(nn.Module):        #定义一个神经网络模型 TheModelClass
        def __init__(self):
            super(TheModelClass, self).__init__()
            self.conv1 = nn.Conv2d(3, 6, 5)
            self.pool = nn.MaxPool2d(2, 2)
            self.conv2 = nn.Conv2d(6, 16, 5)
            self.fc1 = nn.Linear(16 * 5 * 5, 120)
            self.fc2 = nn.Linear(120, 84)
            self.fc3 = nn.Linear(84, 10)
    
        def forward(self, x):
            x = self.pool(F.relu(self.conv1(x)))
            x = self.pool(F.relu(self.conv2(x)))
            x = x.view(-1, 16 * 5 * 5)
            x = F.relu(self.fc1(x))
            x = F.relu(self.fc2(x))
            x = self.fc3(x)
            return x
    
    # 初始化模型
    model = TheModelClass()
    
    # 初始化优化器
    optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
    
    # 打印模型的 state_dict
    print("Model's state_dict:")
    for param_tensor in model.state_dict():        # param_tensor 为参数名称
        print(param_tensor, "\t", model.state_dict()[param_tensor].size())
    
    # 打印优化器的 state_dict
    print("Optimizer's state_dict:")
    for var_name in optimizer.state_dict():
        print(var_name, "\t", optimizer.state_dict()[var_name])
    

    输出结果:

    Model's state_dict:
    conv1.weight torch.Size([6, 3, 5, 5])
    conv1.bias torch.Size([6])
    conv2.weight torch.Size([16, 6, 5, 5])
    conv2.bias torch.Size([16])
    fc1.weight torch.Size([120, 400])
    fc1.bias torch.Size([120])
    fc2.weight torch.Size([84, 120])
    fc2.bias torch.Size([84])
    fc3.weight torch.Size([10, 84])
    fc3.bias torch.Size([10])

    Optimizer's state_dict:
    state {}
    param_groups [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [4675713712, 4675713784, 4675714000, 4675714072, 4675714216, 4675714288, 4675714432, 4675714504, 4675714648, 4675714720]}]

    (3)load提供了很多重载的功能,其可以把在GPU上训练的权重加载到CPU上跑

    torch.load('tensors.pt')
    # 强制所有GPU张量加载到CPU中
    torch.load('tensors.pt', map_location=lambda storage, loc: storage)  #或者model.load_state_dict(torch.load('model.pth', map_location='cpu'))
    # 把所有的张量加载到GPU 1中
    torch.load('tensors.pt', map_location=lambda storage, loc: storage.cuda(1))
    # 把张量从GPU 1 移动到 GPU 0
    torch.load('tensors.pt', map_location={'cuda:1':'cuda:0'})

    上述代码只有在模型在一个GPU上训练时才起作用。如果我在多个GPU上训练模型并且保存它,然后尝试在CPU上加载,会得到错误:KeyError: ‘unexpected key “module.conv1.weight” in state_dict' 如何解决?

    因为此时已经使用模型保存了模型nn.DataParallel,该模型将模型存储在该模型中module,而现在您正试图加载模型DataParallel。您可以nn.DataParallel在网络中暂时添加一个加载目的,也可以加载权重文件,创建一个没有module前缀的新的有序字典,然后加载它。看第(4)点

    (4)通过DataParalle使用多GPU时的保存和加载

    odel=DataParalle(model)
    #保存参数
    torch.save(model.module.state_dict(), 'model.pth')

    由此看出多个GPU时多了一个该模型中module,加载再cpu时,创建一个没有module前缀的新的有序字典,然后加载它。

    补充:一般来说,PyTorch的模型以.pt或者.pth文件格式保存。

    2. 保存加载用于推理的常规Checkpoint/或继续训练**

    checkpoint检查点:不仅保存模型的参数,优化器参数,还有loss,epoch等(相当于一个保存模型的文件夹)

    if (epoch+1) % checkpoint_interval == 0:
    
            checkpoint = {"model_state_dict": net.state_dict(),
                          "optimizer_state_dict": optimizer.state_dict(),
                          "epoch": epoch}
            path_checkpoint = "./checkpoint_{}_epoch.pkl".format(epoch)
            torch.save(checkpoint, path_checkpoint)
    #或者
    #保存
    torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'loss': loss,
                ...
                }, PATH)
    #加载
    model = TheModelClass(*args, **kwargs)
    optimizer = TheOptimizerClass(*args, **kwargs)
    
    checkpoint = torch.load(PATH)
    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    epoch = checkpoint['epoch']
    loss = checkpoint['loss']
    
    model.eval()
    # - 或者 -
    model.train()

    注意:

    在保存用于推理或者继续训练的常规检查点的时候,除了模型的state_dict之外,还必须保存其他参数。保存优化器的state_dict也非常重要,因为它包含了模型在训练时候优化器的缓存和参数。除此之外,还可以保存停止训练时epoch数,最新的模型损失,额外的torch.nn.Embedding层等。

    要保存多个组件,则将它们放到一个字典中,然后使用torch.save()序列化这个字典。一般来说,使用.tar文件格式来保存这些检查点。

    加载各个组件,首先初始化模型和优化器,然后使用torch.load()加载保存的字典,然后可以直接查询字典中的值来获取保存的组件。

    同样,评估模型的时候一定不要忘了调用model.eval()。

    是不是很简单!!以上为个人经验,希望能给大家一个参考,也希望大家多多支持脚本之家。

    您可能感兴趣的文章:
    • pytorch 搭建神经网路的实现
    • 手把手教你实现PyTorch的MNIST数据集
    • Python机器学习之基于Pytorch实现猫狗分类
    • 总结近几年Pytorch基于Imgagenet数据集图像分类模型
    上一篇:Django商城项目注册功能的实现
    下一篇:python 简单的股票基金爬虫
  • 相关文章
  • 

    © 2016-2020 巨人网络通讯 版权所有

    《增值电信业务经营许可证》 苏ICP备15040257号-8

    pytorch模型的保存和加载、checkpoint操作 pytorch,模型,的,保存,和,