精华内容
下载资源
问答
  • pytorch模型的保存与加载保存和加载整个模型仅保存和加载模型参数(推荐使用,需要提前手动构建模型) 保存和加载整个模型 torch.save(model_object, 'model.pkl') model = torch.load('model.pkl') 仅保存和加载模型...

    保存和加载整个模型

    torch.save(model_object, 'model.pkl')
    model = torch.load('model.pkl')
    

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

    torch.save(model_object.state_dict(), 'params.pkl')
    model_object.load_state_dict(torch.load('params.pkl'))
    
    展开全文
  • Pytorch 模型的保存与加载 想先在一部分数据跑出来几个模型,然后选择最优的加载 在另一部分数据上继续跑,加载时遇到了问题,from_pretrained()显示state_dict中的都参数没加载成功,是因为都多了一层encoder: ...

    Pytorch 模型的保存与加载

    想先在一部分数据跑出来几个模型,然后选择最优的加载 在另一部分数据上继续跑,加载时遇到了问题,from_pretrained()显示state_dict中的都参数没加载成功,是因为都多了一层encoder: encoder.encoder.self.attention…所以查找了pytorch保存与加载模型的相关知识,在这里记录一下~

    来源:

    https://pytorch.org/tutorials/beginner/saving_loading_models.html#

    首先需要熟悉三个核心函数:

    1. torch.save: 保存序列化对象到磁盘中。该函数使用pickle做序列化,可以保存模型,张量和字典等对象
    2. torch.load: 使用pickle做反序列化到内存中。
    3. torch.nn.Module.load_state_dict:加载模型的参数字典

    什么是state_dict?

    pytorch里,torch.nn.Module可学习的参数可以通过model.parameters() 得到,而state_dict 就是将这些参数映射到一个python字典里,方便存储和更新。除了卷积层、线性层等参数外,Optimizer(torch.optim)也有state_dict, 包含Optimizer的信息和一些超参数。

    state_dict的例子:

    # Define model
    class TheModelClass(nn.Module):
        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
    
    # Initialize model
    model = TheModelClass()
    
    # Initialize optimizer
    optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
    
    # Print model's state_dict
    print("Model's state_dict:")
    for param_tensor in model.state_dict():
        print(param_tensor, "\t", model.state_dict()[param_tensor].size())
    
    # Print optimizer's state_dict
    print("Optimizer's state_dict:")
    for var_name in optimizer.state_dict():
        print(var_name, "\t", optimizer.state_dict()[var_name])
        
    '''    
    Output:
    
    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]}]
    '''
    

    保存并加载模型

    1. 保存、加载state_dict(推荐)

    保存:

    torch.save(model.state_dict(),PATH)
    

    加载:

    model = TheModelClass(*args, **kwargs)
    model.load_state_dict(torch.load(PATH))
    model.eval()
    

    保存state_dict 会更灵活,方便恢复模型。

    模型保存后的扩展名为.pt.pth

    记住在用模型推理之前要先调用model.eval()设置dropout和batch normalization层,否则会产生不连续的结果

    (这里没调用好像也没产生啥后果)

    注意:

    这里有个很容易出错的地方,经常会误写成model.load_state_dict(PATH),错在哪里呢?

    load_state_dict()参数必须是个字典对象,而不是保存字典的对象。所以先反序列化得到字典对象state_dict= torch.load(PATH), 然后再load

    2. 保存、加载整个模型

    保存:

    torch.save(model,PATH)
    

    加载

    model = torch.load(PATH)
    model.eval()
    

    这种方法最符合我们的直觉,用pickle模块保存整个模型,保存了模型涉及到的具体类的结构,所以当重构模型时可能会出错。

    扩展名同上,为 .pt.pth

    3. 保存、加载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()
    # - or -
    model.train()
    

    当保存checkpoint用于推理或重训练的时候,可能需要epoch/loss等信息,所以需要保存成字典形式再save。

    扩展名一般是.tar

    4. 一个文件里保存多个模型

    保存:

    torch.save({
                'modelA_state_dict': modelA.state_dict(),
                'modelB_state_dict': modelB.state_dict(),
                'optimizerA_state_dict': optimizerA.state_dict(),
                'optimizerB_state_dict': optimizerB.state_dict(),
                ...
                }, PATH)
    

    加载:

    modelA = TheModelAClass(*args, **kwargs)
    modelB = TheModelBClass(*args, **kwargs)
    optimizerA = TheOptimizerAClass(*args, **kwargs)
    optimizerB = TheOptimizerBClass(*args, **kwargs)
    
    checkpoint = torch.load(PATH)
    modelA.load_state_dict(checkpoint['modelA_state_dict'])
    modelB.load_state_dict(checkpoint['modelB_state_dict'])
    optimizerA.load_state_dict(checkpoint['optimizerA_state_dict'])
    optimizerB.load_state_dict(checkpoint['optimizerB_state_dict'])
    
    modelA.eval()
    modelB.eval()
    # - or -
    modelA.train()
    modelB.train()
    

    其实和上个方法相似,不同的模型保存在同一个字典,然后再保存到文件里。

    扩展名同样为.tar

    5. 只使用模型的一部分参数

    保存:

    torch.save(modelA.state_dict(), PATH)
    

    加载:

    modelB = TheModelBClass(*args, **kwargs)
    modelB.load_state_dict(torch.load(PATH), strict=False)
    

    当做迁移学习或加载部分参数来warm_start模型时,参数会不完全匹配,出现 missing keysunexpected keys,这时把load_state_dictstrict设为False即可。

    如果出现需要加载的模型参数名和另一个模型不匹配,可以更改 state_dict的keys,文章开始提到的那个问题,每个参数前面多加了encoder.,所以把每个参数的key去掉了’ encoder.’ 这部分,使用

    # 这里加list是因为,orderDict更改key会报错:
    # OrderedDict mutated during iteration
    # 变成list就可以了
    for param_key in list(state_dict.keys()):
        # 我这里是在最开始多了'encoder.
    	if param_key.find('encoder.') == 0:
            tmp_key = param_key[8:]
            state_dict[tmp_key] = state_dict.pop(param_key)
    

    6. 跨机器保存、加载模型

    在GPU上保存,CPU上加载

    保存:

    torch.save(model.state_dict(), PATH)
    

    加载:

    device = torch.device('cpu')
    model = TheModelClass(*args, **kwargs)
    model.load_state_dict(torch.load(PATH, map_location=device))
    

    当在CPU上加载GPU上训练的模型时,torch.load()需要增加一个map_location参数,重映射到cpu

    在GPU上保存,GPU上加载

    保存:

    torch.save(model.state_dict(), PATH)
    

    加载:

    device = torch.device("cuda")
    model = TheModelClass(*args, **kwargs)
    model.load_state_dict(torch.load(PATH))
    model.to(device)
    # Make sure to call input = input.to(device) on any input tensors that you feed to the model
    

    在GPU上加载一个GPU上训练得到的模型时,需要用model.to(torch.device('cuda'))

    另外,模型的输入也需要使用.to(torch.device('cuda')), 注意,my_tensor.to(device)会返回my_tensor的一个拷贝,所以需要进行覆盖:

    my_tensor = my_tensor.to(torch.device('cuda'))

    CPU上保存,GPU上加载

    保存:

    torch.save(model.state_dict(), PATH)
    

    加载:

    device = torch.device("cuda")
    model = TheModelClass(*args, **kwargs)
    model.load_state_dict(torch.load(PATH, map_location="cuda:0"))  # Choose whatever GPU device number you want
    model.to(device)
    # Make sure to call input = input.to(device) on any input tensors that you feed to the model
    

    在GPU上加载CPU保存的模型时,需要设置map_location参数为 cuda:device_id,然后 model.to(torch.device('cuda')),最后注意所有的输入也需要.to(device)

    保存torch.nn.DataParallel模型

    保存:

    torch.save(model.module.state_dict(), PATH)
    

    加载:

    # Load to whatever device you want
    

    torch.nn.DataParallel设置GPU并行使用,使用model.module.state_dict()保存模型,然后加载的灵活性很强,可以加载到任意device上

    展开全文
  • Pytorch模型的保存与加载

    万次阅读 2018-08-30 16:58:20
    前言 在使用Pytorch训练模型的时候,经常会有在GPU上保存模型...模型的保存与加载 例如我们创建了一个模型: model = MyVggNet() 如果使用多GPU训练,我们需要使用这行代码: model = nn.DataPa...

    前言

    在使用Pytorch训练模型的时候,经常会有在GPU上保存模型然后再CPU上运行的需求,在实验的过程中发现在多GPU上训练的Pytorch模型是不能在CPU上直接运行的,几次遇到了这种问题,这里研究和记录一下。

    模型的保存与加载

    例如我们创建了一个模型:

    model = MyVggNet()
    

    如果使用多GPU训练,我们需要使用这行代码:

    model = nn.DataParallel(model).cuda()
    

    执行这个代码之后,model就不在是我们原来的模型,而是相当于在我们原来的模型外面加了一层支持GPU运行的外壳,这时候真正的模型对象为:real_model = model.module, 所以我们在保存模型的时候注意,如果保存的时候是否带有这层加的外壳,如果保存的时候带有的话,加载的时候也是带有的,如果保存的是真实的模型,加载的也是真是的模型。这里我建议保存真是的模型,因为加了module壳的模型在CPU上是不能运行的。
    Pytorch有多种保存模型的方式,使用哪种进行保存,就要使用对应的加载方式。保存的时候模型的后缀名是无所谓的。
    Pytorch官方的加载和保存模型的方式有两种:

    1. 保存和加载整个模型。这种方式再重新加载的时候不需要自定义网络结构,保存时已经把网络结构保存了下来,比较死板不能调整网络结构。
    torch.save(model_object, 'model.pkl')
    model = torch.load('model.pkl')
    
    1. 仅保存和加载模型参数(推荐使用)。这种方式再重新加载的时候需要自己定义网络,并且其中的参数名称与结构要与保存的模型中的一致(可以是部分网络,比如只使用VGG的前几层),相对灵活,便于对网络进行修改。
    torch.save(model_object.state_dict(), 'params.pkl')
    model_object.load_state_dict(torch.load('params.pkl'))
    

    模型保存与加载对应方式

    1. 第一种方式

    保存使用:

    real_model = model.module
    torch.save(real_model.state_dict(),os.path.join(args.save_path,"cos_mnist_"+str(epoch+1)+"_weight.pth"))
    

    cpu上加载使用:

    args.weight=checkpoint/cos_mnist_10_weight.pth
    map_location = lambda storage, loc: storage
    model.load_state_dict(torch.load(args.weight,map_location=map_location))
    

    2. 第二种方式

    保存使用:

    real_model = model.module
    save_model(real_model, os.path.join(args.save_path,"cos_mnist_"+str(epoch+1)+"_weight_cpu.pth"))
    # 自定义的函数
    def save_model(model,filename):
        state = model.state_dict()
        for key in state: state[key] = state[key].clone().cpu()
        torch.save(state, filename)
    

    cpu上加载使用:

    args.weight=checkpoint/cos_mnist_10_weight_cpu.pth
    model.load_state_dict(torch.load(args.weight))
    

    3. 第三种方式

    保存使用:

    real_model = model.module
    torch.save(real_model, os.path.join(args.save_path,"cos_mnist_"+str(epoch+1)+"_whole.pth"))
    

    cpu上加载使用:

    args.weight=checkpoint/cos_mnist_10_whole.pth
    map_location = lambda storage, loc: storage
    model = torch.load(args.weight,map_location=map_location)
    

    参考文献

    1. pytorch学习笔记(五):保存和加载模型
    2. Pytorch之GPU模型加载在CPU上
    3. PyTorch使用cpu调用gpu训练的模型
    展开全文
  • 将pytorch模型保存为c++程序可直接访问的模型: ... pytorch模型的序列化: 官方指南: https://pytorch.org/docs/stable/notes/serialization.html#best-practices ...

    将pytorch模型保存为c++程序可直接访问的模型:

    https://pytorch.org/tutorials/advanced/cpp_export.html

     

    pytorch模型的序列化:

    官方指南:

    https://pytorch.org/docs/stable/notes/serialization.html#best-practices

    展开全文

空空如也

空空如也

1 2 3 4 5 ... 19
收藏数 362
精华内容 144
关键字:

pytorch模型的保存与加载