pytorch 模型保存的完整例子+pytorch 模型保存只保存可训练参数吗?是(+解决方案)

本文介绍如何在PyTorch中保存和加载模型,包括只保存可训练参数的方法,并演示了一个具体的例子。同时探讨了如何处理非训练参数以及保存优化器状态。
部署运行你感兴趣的模型镜像

        测试使用的是一个liner model,还有更多的问题。pytorch 模型保存只保存可训练参数吗?

save模型

# 导入包
import glob
import os

import torch
import matplotlib.pyplot as plt
import random #用于数据迭代器生成随机数据

# 生成数据集 x1类别0,x2类别1
n_data = torch.ones(50, 2)  # 数据的基本形态
x1 = torch.normal(2 * n_data, 1)  # shape=(50, 2)
y1 = torch.zeros(50)  # 类型0 shape=(50, 1)
x2 = torch.normal(-2 * n_data, 1)  # shape=(50, 2)
y2 = torch.ones(50)  # 类型1 shape=(50, 1)
# 注意 x, y 数据的数据形式一定要像下面一样(torch.cat是合并数据)
x = torch.cat((x1, x2), 0).type(torch.FloatTensor)
y = torch.cat((y1, y2), 0).type(torch.FloatTensor)

# 数据集可视化
plt.scatter(x.data.numpy()[:, 0], x.data.numpy()[:, 1], c=y.data.numpy(), s=100, lw=0, cmap='RdYlGn')
plt.show()

# 数据读取:
def data_iter(batch_size, x, y):
    num_examples = len(x)
    indices = list(range(num_examples))
    random.shuffle(indices)  # 样本的读取顺序是随机的
    for i in range(0, num_examples, batch_size):
        j = torch.LongTensor(indices[i: min(i + batch_size, num_examples)]) #最后一次可能不足一个batch
        yield  x.index_select(0, j), y.index_select(0, j)

#############################################################################################################
def saver(model_state_dict, optimizer_state_dict, model_path, epoch, max_to_save=30):
    total_models = glob.glob(model_path + '*')
    if len(total_models) >= max_to_save:
        total_models.sort()
        os.remove(total_models[0])

    state_dict = {}
    state_dict["model_state_dict"] = model_state_dict
    state_dict["optimizer_state_dict"] = optimizer_state_dict

    torch.save(state_dict, model_path + 'h' + str(epoch))
    print('models {} save successfully!'.format(model_path + 'hahaha' + str(epoch)))



################################################################################################################

import torch.nn as nn
import torch.optim as optim



class net(nn.Module):
    def __init__(self, **kwargs):
        super(net, self).__init__(**kwargs)
        self.net = nn.Sequential(nn.Linear(2, 1), nn.ReLU())

    def forward(self, x):
        return self.net(x)

def loss(y_hat, y):
    return (y_hat - y.view(y_hat.size())) ** 2 / 2



def accuracy(y_hat, y):  #@save
    """计算预测正确的数量。"""
    cmp = y_hat.type(y.dtype) > 0.5 # 大于0.5类别1
    result=cmp.type(y.dtype)
    acc = 1-float(((result-y).sum())/ len(y))
    return acc;

lr = 0.03
num_epochs = 3 # 迭代次数
batch_size = 10 # 批量大小
model = net()
params =  list(model.parameters())
optimizer = torch.optim.Adam(params, 1e-4)

for epoch in range(num_epochs):
    for X, y_train in data_iter(batch_size, x, y):
        optimizer.zero_grad()
        l = loss(model(X), y_train).sum()  # l是有关小批量X和y的损失
        l.backward(retain_graph=True)
        optimizer.step()
        print(l)
    saver(model.state_dict(), optimizer.state_dict(), "./", epoch + 1,  max_to_save=100)



load模型

# 导入包
import glob
import os

import torch
import matplotlib.pyplot as plt
import random #用于数据迭代器生成随机数据

# 生成数据集 x1类别0,x2类别1
n_data = torch.ones(50, 2)  # 数据的基本形态
x1 = torch.normal(2 * n_data, 1)  # shape=(50, 2)
y1 = torch.zeros(50)  # 类型0 shape=(50, 1)
x2 = torch.normal(-2 * n_data, 1)  # shape=(50, 2)
y2 = torch.ones(50)  # 类型1 shape=(50, 1)
# 注意 x, y 数据的数据形式一定要像下面一样(torch.cat是合并数据)
x = torch.cat((x1, x2), 0).type(torch.FloatTensor)
y = torch.cat((y1, y2), 0).type(torch.FloatTensor)

# 数据集可视化
plt.scatter(x.data.numpy()[:, 0], x.data.numpy()[:, 1], c=y.data.numpy(), s=100, lw=0, cmap='RdYlGn')
plt.show()

# 数据读取:
def data_iter(batch_size, x, y):
    num_examples = len(x)
    indices = list(range(num_examples))
    random.shuffle(indices)  # 样本的读取顺序是随机的
    for i in range(0, num_examples, batch_size):
        j = torch.LongTensor(indices[i: min(i + batch_size, num_examples)]) #最后一次可能不足一个batch
        yield  x.index_select(0, j), y.index_select(0, j)

#############################################################################################################
def saver(model_state_dict, optimizer_state_dict, model_path, epoch, max_to_save=30):
    total_models = glob.glob(model_path + '*')
    if len(total_models) >= max_to_save:
        total_models.sort()
        os.remove(total_models[0])

    state_dict = {}
    state_dict["model_state_dict"] = model_state_dict
    state_dict["optimizer_state_dict"] = optimizer_state_dict

    torch.save(state_dict, model_path + 'h' + str(epoch))
    print('models {} save successfully!'.format(model_path + 'hahaha' + str(epoch)))



################################################################################################################

import torch.nn as nn
import torch.optim as optim



class net(nn.Module):
    def __init__(self, **kwargs):
        super(net, self).__init__(**kwargs)
        self.net = nn.Sequential(nn.Linear(2, 1), nn.ReLU())

    def forward(self, x):
        return self.net(x)

def loss(y_hat, y):
    return (y_hat - y.view(y_hat.size())) ** 2 / 2



def accuracy(y_hat, y):  #@save
    """计算预测正确的数量。"""
    cmp = y_hat.type(y.dtype) > 0.5 # 大于0.5类别1
    result=cmp.type(y.dtype)
    acc = 1-float(((result-y).sum())/ len(y))
    return acc;

lr = 0.03
num_epochs = 3 # 迭代次数
batch_size = 10 # 批量大小
model = net()
params =  list(model.parameters())
optimizer = torch.optim.Adam(params, 1e-4)

# for epoch in range(num_epochs):
#     for X, y_train in data_iter(batch_size, x, y):
#         optimizer.zero_grad()
#         l = loss(model(X), y_train).sum()  # l是有关小批量X和y的损失
#         l.backward(retain_graph=True)
#         optimizer.step()
#         print(l)
#     saver(model.state_dict(), optimizer.state_dict(), "./", epoch + 1,  max_to_save=100)




def loader(model_path):
    state_dict = torch.load(model_path)
    model_state_dict = state_dict["model_state_dict"]
    optimizer_state_dict = state_dict["optimizer_state_dict"]
    return model_state_dict, optimizer_state_dict

model_state_dict, optimizer_state_dict = loader("h1")
model.load_state_dict(model_state_dict)
optimizer.load_state_dict(optimizer_state_dict)

print('pretrained models loaded!')

pytorch 模型保存只保存可训练参数吗?是

class net(nn.Module):
    def __init__(self, **kwargs):
        super(net, self).__init__(**kwargs)
        self.net = nn.Sequential(nn.Linear(2, 1), nn.ReLU())
        self.notrain= torch.rand((64, 64), dtype=torch.float)

    def forward(self, x):
        return self.net(x)

在这里插入图片描述

解决方案

  • 直接更改.data
class net(nn.Module):
    def __init__(self, **kwargs):
        super(net, self).__init__(**kwargs)
        self.net = nn.Sequential(nn.Linear(2, 1), nn.ReLU())
        # self.notrain = torch.rand((64, 64), dtype=torch.float)
        self.notrain = torch.nn.Parameter(torch.ones(64, 64))

    def forward(self, x):
        return self.net(x)
for epoch in range(num_epochs):
    for X, y_train in data_iter(batch_size, x, y):
        optimizer.zero_grad()
        l = loss(model(X), y_train).sum()  # l是有关小批量X和y的损失
        l.backward(retain_graph=True)
        optimizer.step()
        print(l)
        model.notrain.data = model.notrain.data+2
    saver(model.state_dict(), optimizer.state_dict(), "./", epoch + 1,  max_to_save=100)

TypeError: cannot assign ‘torch.cuda.FloatTensor’ as parameter ‘***’ (torch.nn.Parameter or None expected)

  • self.weight = self.weight.detach()会报以上的错误,可以考虑使用
  1. 在网络传播中detach(这种方法一般效率低)
  2. 推荐注册为buffer,或者直接self.weight = torch.nn.Parameter(Tensor data, requires_grad = False)
  3. model.*** = torch.nn.Parameter(torch.load("./SAVEPE.pt"))

参考与更多

PyTorch DataLoader的bug :随机mask或者对数据的随机挑选产生的bug

您可能感兴趣的与本文相关的镜像

PyTorch 2.5

PyTorch 2.5

PyTorch
Cuda

PyTorch 是一个开源的 Python 机器学习库,基于 Torch 库,底层由 C++ 实现,应用于人工智能领域,如计算机视觉和自然语言处理

<think>好的,我现在需要解决用户在加载PyTorch模型遇到的KeyError问题。根据用户提供的引用内容,特别是引用[1]中的错误信息,用户可能在加载预训练模型的权重遇到了键不匹配的问题。首先,我应该回忆一下常见的导致KeyError的原因。 首先,模型结构不匹配是一个常见原因。例如,用户可能修改了模型的某些层名称,导致保存的权重中的键名与当前模型的层名不一致。比如,原模型可能有`bert.embeddings.word_embeddings.weight`,而新模型可能改名或删除了该层,导致加载找不到对应的键,从而引发KeyError。 其次,用户可能使用了不同的预训练模型配置。例如,原本的模型是基于BERT-base,而用户现在使用的是BERT-large,这样层数和参数结构不同,自然键名也会不同。这种情况下,加载权重就会出现键不匹配。 另外,用户可能在保存模型没有正确保存整个模型,而只保存了状态字典的一部分,或者在加载使用了错误的模型类。比如,原本用`BertForSequenceClassification`保存模型,现在尝试用`BertForTokenClassification`加载,这会导致部分键缺失。 接下来,我需要考虑解决方案。根据引用[1]中的例子,用户遇到了形状不匹配的问题,但这里的问题更偏向于键名不存在。可能的解决步骤包括: 1. **检查模型结构是否一致**:确保当前定义的模型保存权重模型结构完全相同。可以通过打印模型的`state_dict`的键来比较。 2. **使用严格加载模式**:PyTorch的`load_state_dict`函数有一个`strict`参数,默认为True。如果设为False,可以忽略不匹配的键,但需要谨慎,因为这可能导致部分参数未加载。 3. **手动调整键名**:如果键名有规律的变化(比如前缀不同),可以编写脚本批量修改加载的权重字典中的键名,使其匹配当前模型。 4. **检查模型类和版本**:确认使用的PyTorch版本和模型定义库(如HuggingFace Transformers)的版本是否与训练一致。版本更新可能导致接口变化。 5. **部分加载参数**:如果只有部分层不匹配,可以分别加载匹配的部分参数,而忽略不存在的键。例如,遍历`state_dict`,只加载当前模型存在的键。 此外,用户需要确认错误信息中的具体KeyError内容,例如缺失的键名是什么,这有助于定位问题所在。例如,如果错误是`KeyError: 'bert.embeddings.position_embeddings.weight'`,则可能该层在当前模型中被移除或改名。 参考引用[2]提到的模型实现(如FFDNet),可能需要关注自定义层是否导致键名变化。同,引用[3]中的类型错误虽然不直接相关,但提醒我需要确保数据类型的一致性,不过KeyError更多是结构问题。 总结可能的解决步骤后,我需要以清晰的中文逐步指导用户检查这些方面,并提供示例代码帮助理解如何实施解决方案,比如修改键名或设置`strict=False`。同,生成相关问题以帮助用户深入理解类似问题。</think>### PyTorch模型加载KeyError解决方案 #### 问题分析 加载PyTorch模型出现`KeyError`,通常由以下原因导致: 1. **模型结构不匹配**:当前模型定义与保存的权重文件结构不一致(如层名称或参数形状不同)[^1] 2. **权重文件损坏**:文件保存完整保存了非`state_dict`对象 3. **版本兼容性问题**:PyTorch或第三方库版本升级导致接口变化 --- #### 解决方案分步说明 ##### 步骤1:验证权重文件完整性 ```python import torch try: checkpoint = torch.load("model.pth") print("Keys in checkpoint:", checkpoint.keys()) # 查看权重文件结构 except Exception as e: print("文件加载失败:", str(e)) ``` - 正常应包含`state_dict`键 - 若出现`EOFError`,说明文件损坏 ##### 步骤2:检查模型结构一致性 ```python from models import YourModelClass # 创建新模型实例 model = YourModelClass() # 打印当前模型参数名称 print("Current model keys:", [k for k in model.state_dict().keys()][:3]) # 示例:['bert.embeddings.word_embeddings.weight', ...] # 打印权重文件参数名称 checkpoint = torch.load("model.pth") print("Checkpoint keys:", [k for k in checkpoint['state_dict'].keys()][:3]) ``` - 比较输出结果中的键名差异 ##### 步骤3:使用严格加载模式(推荐方案) ```python # 方案1:严格模式加载(需要完全匹配) model.load_state_dict(checkpoint['state_dict'], strict=True) # 默认模式 # 方案2:非严格模式加载(允许部分加载) missing_keys, unexpected_keys = model.load_state_dict( checkpoint['state_dict'], strict=False ) print("缺失的键:", missing_keys) print("意外的键:", unexpected_keys) ``` - 若出现`missing_keys`,说明当前模型缺少某些层定义 - 若出现`unexpected_keys`,说明权重文件包含多余参数 ##### 步骤4:键名映射修正(高级操作) 当键名存在规律性差异(如多出前缀`module.`),可手动修正: ```python from collections import OrderedDict # 移除所有键名前缀'module.' new_state_dict = OrderedDict() for k, v in checkpoint['state_dict'].items(): name = k.replace('module.', '') # 修改键名规则 new_state_dict[name] = v model.load_state_dict(new_state_dict) ``` ##### 步骤5:版本兼容性处理 ```bash # 查看训练使用的版本 pip freeze > requirements.txt # 回退到指定版本 pip install torch==1.9.0+cu111 transformers==4.12.3 ``` --- #### 典型错误场景 1. **微调后层名变化** - 现象:`KeyError: 'classifier.weight'` - 原因:原始模型有`fc`层,微改后命名为`classifier` - 方案:使用`strict=False`或修改层名 2. **多GPU训练差异** - 现象:键名前缀包含`module.` - 方案:使用步骤4的键名修正代码 3. **第三方库更新** - 现象:HuggingFace Transformers库升级后层名改变 - 方案:固定库版本或手动调整键名映射 ---
评论
成就一亿技术人!
拼手气红包6.0元
还能输入1000个字符
 
红包 添加红包
表情包 插入表情
 条评论被折叠 查看
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值