深入理解 PyTorch .pth 文件和 Python pickle 模块:功能、应用及实际示例

深入理解 PyTorch .pth 文件和 Python pickle 模块:功能、应用及实际示例

在深入理解Python的pickle模块和PyTorch的.pth文件,以及pickle.pth文件中的应用前,我们首先需要掌握序列化和反序列化的基本概念。

序列化和反序列化

序列化 是指将程序中的对象转换为一个字节序列的过程,这样就可以将其存储到磁盘上或通过网络传输到其他位置。反序列化是序列化的逆过程,即将字节序列恢复为原始对象。这两个过程是数据持久化和远程计算通信的基础。

Python的pickle模块

pickle是Python的标准库之一,提供了一个简单的方法用于序列化和反序列化Python对象结构。任何Python对象都可以通过pickle进行序列化,只要它们是pickle支持的类型。

核心功能

  • pickle.dump(obj, file) :将对象obj序列化并写入到文件对象file中。
  • pickle.load(file) :从文件对象file中读取序列化的对象并反序列化。
  • pickle.dumps(obj) :将对象obj序列化为一个字节对象,不写入文件。
  • pickle.loads(bytes_object) :将字节对象bytes_object反序列化为一个Python对象。

pickle的序列化过程不仅包括对象当前的状态(例如,数字,字符串,或复杂对象的集合),也包括对象的类型信息和结构。

PyTorch的.pth文件

在PyTorch中,.pth文件扩展通常用于保存模型的权重或整个模型。这些文件通过使用torch.save()函数创建,它内部使用pickle来序列化对象。.pth文件通常包含模型的状态字典(state_dict),这是一个从每个层的参数名称映射到其张量值的字典。

核心用途

  • 模型持久化:保存训练后的模型状态,以便将来可以重新加载和使用模型,不需要重新训练。
  • 模型迁移:将训练好的模型参数迁移到新的模型结构或平台上。

pickle.pth文件中的应用

当使用torch.save()来保存一个PyTorch模型或张量时,pickle用于将对象和它的层次结构转换为一个字节流,然后这个字节流被写入到一个.pth文件中。在加载模型时,torch.load()使用pickle来反序列化这个字节流,重建模型或张量。

示例

python 复制代码
import torch
import torchvision.models as models

# 实例化一个预训练的ResNet模型
model = models.resnet18(pretrained=True)

# 保存模型状态字典
torch.save(model.state_dict(), 'model_weights.pth')

# 加载模型状态字典
loaded_state_dict = torch.load('model_weights.pth')
new_model = models.resnet18(pretrained=False)
new_model.load_state_dict(loaded_state_dict)

# 打印以验证加载
print(new_model)

在这个示例中,torch.save()内部使用pickle来序列化model.state_dict(),并将其保存为model_weights.pth。然后,我们使用torch.load()来加载这个.pth文件,其中pickle负责反序列化文件内容,并恢复为Python对象(在这种情况下是模型的状态字典)。最后,状态字典被用于初始化一个新的模型实例。

通过这种方式,pickle在PyTorch的模型保存和加载过程中扮演了核心角色,使得模型的状态可以在不同的计算环境中被迁移和复用。

相关推荐
程序媛kelly1 分钟前
如何打开 .md / .ipynb 文件?Markdown 与 Jupyter Notebook 本地预览全攻略
ide·python·jupyter
wjcroom2 分钟前
时空和电子7-泡力模型含罗量
人工智能·算法·机器学习
Codebee4 分钟前
从 Agent Harness 到 Scenario Loop
人工智能
MatrixOrigin6 分钟前
MatrixOne Git4Data 技术详解(三):MatrixOne 架构及 Git4Data 原理解析,快照、Diff、Merge 凭什么这么快
数据库·人工智能·数据平台·矩阵起源·数据底座
KaMeidebaby6 分钟前
卡梅德生物技术快报 | Fab 合成文库构建与抗体筛选实验流程及数据解析
人工智能·python·tcp/ip·算法·机器学习
天风之翼6 分钟前
AI 全栈开发实战(8):前端开发(二)——流式对话界面与 Markdown 渲染
人工智能
培培说证7 分钟前
大数据、人工智能、计算机、软件工程,到底怎么选?
大数据·人工智能·软件工程
装不满的克莱因瓶8 分钟前
掌握3D CNN模型结构——从时空特征建模到视频理解与医学影像核心架构
人工智能·pytorch·python·深度学习·神经网络·3d·cnn
AINative软件工程9 分钟前
LLM 应用的 Schema 演进工程:structured output 字段改了,下游为什么炸了?
后端·python·架构
金融小师妹10 分钟前
基于AI事件驱动模型与验证溢价框架的市场分析:从预期交易到事实验证,原油与黄金面临关键定价重构
大数据·人工智能·算法·均值算法·线性回归