pytorch 模型保存与加载

秒速五厘米 2023-10-03 11:11 2阅读 0赞

一、模型保存有两种形式：保存整体模型（包括模型结构和参数）、只保存模型参数

import torch
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # 保存整体模型
    output_dir = 'checkpoint.ckp'
    model_to_save = model.module if hasattr(model, "module") else model
    torch.save(model_to_save, output_dir)
    # 加载整体模型
    model = torch.load(output_dir, map_location=device)
    
    # =========================================================================
    # 只保存模型参数 state_dict
    torch.save(model.state_dict(), output_dir)
    # 只加载模型参数 state_dict
    model.load_state_dict(torch.load(output_dir, map_location=device))

二、当训练的代码中使用了“**torch.nn.DataParallel()**”，这个命令是将网络在多块gpu中进行训练然后合并，这时采用上面“只保存模型参数”的方式时，保存的参数key中会在最前面多一个module.

解决方式有三个：

1.加载模型时去掉key中的module.

model.load_state_dict({k.replace('module.',''):v for k,v in torch.load('checkpoint.pth').items()})

2.加载时也用“**torch.nn.DataParallel()**”

if cuda:
            g_model = torch.nn.DataParallel(g_model)
            cudnn.benchmark = True
            g_model = g_model.cuda()
    
        if os.path.exists(model_path):
            print('Loading weights into state dict...')
            model_dict = g_model.state_dict()
            pretrained_dict = torch.load(model_path, map_location=device)
            g_model.load_state_dict(pretrained_dict)
            print('Finished!')

3.只有一个GPU就没有必要使用“**torch.nn.DataParallel()**”了

发表评论取消回复

表情：

评论列表（有 0 条评论，2人围观）

还没有评论，来说两句吧...

相关阅读

相关 PyTorch模型保存与加载

. `torch.save`：保存序列化的对象到磁盘，使用了Python的pickle进行序列化，模型、张量、所有对象的字典。 2. `torch.load`：使用了...

Bertha 。/ 2024年04月17日 16:16/ 0 赞/ 70 阅读

相关 Pytorch模型参数的保存和加载介绍

![be721a1d89a7939f331c03c8ac2888e1.jpeg][] pytorch中state\_dict()和load\_state\_dict()函数配

约定不等于承诺〃/ 2024年03月22日 20:40/ 0 赞/ 80 阅读

相关 PyTorch模型保存与加载

PyTorch模型保存与加载在利用PyTorch构建深度学习模型时，模型的保存和加载是非常重要的一步。这不仅可以保证我们的模型得以长期保存和重复使用，还可以方便我们在不同的

清疚/ 2024年03月17日 23:41/ 0 赞/ 79 阅读

相关 pytorch 模型保存与加载

一、模型保存有两种形式：保存整体模型（包括模型结构和参数）、只保存模型参数 import torch device = torch.device('

秒速五厘米/ 2023年10月03日 11:11/ 0 赞/ 2 阅读

相关 pytorch保存模型pth_PyTorch之保存加载模型

前提 SAVING AND LOADING MODELS 当提到保存和加载模型时，有三个核心功能需要熟悉： 1.torch.save：将序列化的对象保存到disk。这个函

Bertha 。/ 2022年12月31日 15:30/ 0 赞/ 293 阅读

相关 sklearn模型保存与加载

sklearn模型保存与加载 sklearn模型的保存和加载API 线性回归的模型保存加载案例保存模型 sklearn模型的保存和加载API

我会带着你远行/ 2022年12月29日 14:00/ 0 赞/ 244 阅读

相关 pytorch保存与加载模型来测试或继续训练

目录摘要 state\_dict 恢复训练实例加载部分预训练模型保存 & 加载模型来inference 保存/加载sta

蔚落/ 2022年12月27日 09:22/ 0 赞/ 237 阅读

相关 pytorch保存模型、加载模型的两种方式

第一种保存整个模型文件保存和加载整个模型 torch.save(model_object, 'model.pth') model = torc

刺骨的言语ヽ痛彻心扉/ 2022年12月25日 10:50/ 0 赞/ 212 阅读

相关 Pytorch之模型加载/保存

pytorch保存模型有两种方法： 1. 保存整个模型 (结构+参数) 2. 只保存参数（官方推荐）两者都是用`torch.save(obj, dir)`实现，这个函

た入场券/ 2022年11月27日 00:51/ 0 赞/ 263 阅读

相关 Tensorflow模型保存与加载

Tensorflow模型保存与加载 import os os.environ['TF_CPP_MIN_LOG_LEVEL']='2'

小咪咪/ 2022年09月08日 00:09/ 0 赞/ 343 阅读