python

关注公众号 jb51net

关闭
首页 > 脚本专栏 > python > PyTorch实现.pt与.bin的转换

在PyTorch中实现.pt模型文件与.bin二进制格式的高效转换

作者:Asia-Lee

文章介绍了.pt和.bin两种模型保存格式的区别,包括格式规范、用途、兼容性等,并提供了转换方法和典型应用场景,特别强调了在不同场景下选择合适格式的重要性,以及注意事项和示例,如HuggingFace模型的转换,需要的朋友可以参考下

1. 格式区别

特性.pt(PyTorch格式).bin(通用二进制格式)
开发者/框架PyTorch 专用通用格式,无框架限制
存储内容可保存完整的模型(参数+架构+优化器状态+其他元数据)通常仅保存模型权重(参数)
文件大小较大(包含额外信息)较小(仅参数)
加载方式需用 torch.load() 和对应模型类需手动加载到匹配的模型结构中
兼容性依赖 PyTorch 版本跨框架兼容(如 ONNX、TensorFlow)
典型用途PyTorch 训练/推理跨平台部署、参数共享

2. 转换方法

2.1 .pt → .bin(提取权重)

import torch
# 加载原始 .pt 文件(假设为模型权重)
model = torch.load("model.pt")  # 可能是完整的模型或仅参数
# 提取权重字典(若保存的是完整模型,需 model.state_dict())
if isinstance(model, torch.nn.Module):
    state_dict = model.state_dict()
else:
    state_dict = model
# 保存为 .bin 文件
torch.save(state_dict, "model.bin")

2.2 .bin → .pt(重建完整模型)

import torch
from model_definition import MyModel  # 导入模型定义
# 初始化模型
model = MyModel()
# 加载 .bin 文件中的参数
state_dict = torch.load("model.bin")
model.load_state_dict(state_dict)
# 保存为完整 .pt 文件(可选)
torch.save(model, "full_model.pt")

3. 典型场景

3.1 跨框架部署

3.2 参数共享

3.3 轻量化部署

4. 注意事项

  1. 模型架构一致性
    • .bin 文件不包含模型结构,加载时需确保目标模型与参数完全匹配。
    • 若参数名称或形状不匹配,会导致加载失败(如 Missing key(s) in state_dict)。
  2. 版本兼容性
    • PyTorch 的 .pt 文件可能因版本不同导致兼容性问题,而 .bin 作为纯权重文件更易跨版本使用。
  3. 安全性
    • .pt 文件可能包含恶意代码(因其可序列化任意对象),加载来源不明的文件需谨慎。
    • .bin 仅参数,安全性相对更高。

5. 示例:Hugging Face 模型转换

Hugging Face 模型通常提供两种格式:

转换代码

from transformers import AutoModel
# 加载 PyTorch 模型
model = AutoModel.from_pretrained("bert-base-uncased")
# 保存为 .bin 文件(仅权重)
model.save_pretrained("save_dir", save_format="pt")  # 生成 pytorch_model.bin

6. 总结

以上就是在PyTorch中实现.pt模型文件与.bin二进制格式的高效转换的详细内容,更多关于PyTorch实现.pt与.bin的转换的资料请关注脚本之家其它相关文章!

您可能感兴趣的文章:
阅读全文