Pytorch基础:torch.load_state_dict()方法在加载时不会检查类型

相关阅读

Pytorch基础https://blog.csdn.net/weixin_45791458/category_12457644.html?spm=1001.2014.3001.5482


笔者在使用torch.nn.module的load_state_dict中出现了一个问题,一个被注册的张量在加载后居然没有变化,一开始以为是加载出现了问题,但发现其他参数加载成功,思索后发现是注册的张量的类型是整型而checkpoint中保存为浮点数类型,恰好注册时的默认值给的是0,而checkpoint中的浮点数又在0到1之间,因此出现了这个令人困惑的bug。

下面首先复现这个bug。

复制代码
import torch
import torch.nn as nn

# 定义一个简单的线性模型,参数类型为整数
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.register_buffer('test', torch.tensor(0)) # 注册一个整型张量

# 创建一个简单模型实例
model = SimpleModel()

# 创建一个浮点数作为参数
float_parameter = torch.tensor(0.6)

# 将注册名指向另一个浮点型张量
model.test = float_parameter

# 保存模型
torch.save(model.state_dict(), 'model.pth')

# 直接使用原模型加载
checkpoint = torch.load('model.pth')
model.load_state_dict(checkpoint)

# 打印加载后的参数
print(model.test)

# 直接使用新模型加载
model_1 = SimpleModel()
model_1.load_state_dict(checkpoint)

# 打印加载后的参数
print(model_1.test)

输出:
tensor(0.6000)
tensor(0)

可以看到,当模型中注册的名字(test),指向了一个类型不符的张量后,并不会导致浮点型张量被截断为整型,这是因为此处是直接使用赋值号=,使名字指向了另一个张量。

但使用load_state_dict()方法与使用赋值号是不同的,load_state_dict()方法的实现中,调用了_load_from_state_dict()方法,其中调用了copy_()方法,进行了原位(in-place)数据替换,这可能会进行截断,下面是原位替换的一个例子。

python 复制代码
import torch

# 创建两个张量
a = torch.tensor([[1, 2], [3, 4]])
b = torch.tensor([[5.1, 6.1], [7.1, 8.1]])

# 查看张量对象的id
print(id(a))
print(id(b))

# 查看底层存储的内存地址
print(a.storage().data_ptr())
print(b.storage().data_ptr())

# 将张量 b 中的值复制到张量 a 中
a.copy_(b)

# 打印复制后的结果
print(a)

# 查看张量对象的id
print(id(a))
print(id(b))

# 查看底层存储的内存地址
print(a.storage().data_ptr())
print(b.storage().data_ptr())
python 复制代码
输出:
2604425272672
2604426953808  
2604511348096  
2602930352832  
tensor([[5, 6],
        [7, 8]])
2604425272672
2604426953808
2604511348096
2602930352832

在保存了模型的状态字典后,使用load_state_dict()方法加载后,也不会有任何截断问题,因为对于原模型而言,名字test指向的是一个浮点型张量,此时原位替换,类型吻合。但是对于一个新的模型,此时的test指向的是一个整型张量,此时原位替换,会发生截断。

因此,在注册一个张量时,需要确保其在注册时和保存时的类型吻合,此处除了指形状,还有类型,否则可能会出现意想不到的bug。

相关推荐
nju_spy1 分钟前
周志华《机器学习导论》第8章 集成学习 Ensemble Learning
人工智能·随机森林·机器学习·集成学习·boosting·bagging·南京大学
心情好的小球藻21 分钟前
Python应用进阶DAY9--类型注解Type Hinting
开发语言·python
都叫我大帅哥23 分钟前
LangChain加载HTML内容全攻略:从入门到精通
python·langchain
静心问道25 分钟前
TrOCR: 基于Transformer的光学字符识别方法,使用预训练模型
人工智能·深度学习·transformer·多模态
说私域27 分钟前
基于开源AI大模型、AI智能名片与S2B2C商城小程序源码的用户价值引导与核心用户沉淀策略研究
人工智能·开源
亲持红叶28 分钟前
GLU 变种:ReGLU 、 GEGLU 、 SwiGLU
人工智能·深度学习·神经网络·激活函数
说私域28 分钟前
线上协同办公时代:以开源AI大模型等工具培养网感,拥抱职业变革
人工智能·开源
群联云防护小杜30 分钟前
深度隐匿源IP:高防+群联AI云防护防绕过实战
运维·服务器·前端·网络·人工智能·网络协议·tcp/ip
惜.己33 分钟前
使用python读取json数据,简单的处理成元组数组
开发语言·python·测试工具·json
摘星编程35 分钟前
构建智能客服Agent:从需求分析到生产部署
人工智能·需求分析·智能客服·agent开发·生产部署