📝 引言

在深度学习领域,模型文件格式的选择直接影响到模型的存储、加载效率、安全性和跨平台兼容性。本文将全面解析主流的模型文件格式,帮助读者在不同应用场景下做出最优选择。


💾 主流模型文件格式详解

1. .safetensors - 新一代安全高效的模型格式

核心特性
  • 开发方:由 Hugging Face 团队开发
  • 设计目标:安全性、高效加载、跨平台兼容
  • 文件内容:仅包含模型的权重参数,不包含可执行代码
  • 存储结构:优化的二进制格式,支持内存映射
主要优势
  • 安全性:完全避免了代码注入风险
  • 加载速度
    • 通过内存映射技术,显著提升大模型加载速度
    • 支持零拷贝加载(Zero-copy loading),直接从磁盘映射到内存,无需额外的内存拷贝操作
    • 特别适合加载大型模型,可显著减少内存使用和加载时间
  • 灵活访问:支持选择性加载特定层的权重
  • 内存效率:优化的文件布局,减少内存占用
  • 适用场景:Hugging Face 生态、快速部署

2. .ckpt - PyTorch Lightning 的标准格式

核心特性
  • 适用框架:PyTorch Lightning 生态
  • 存储范围:完整训练状态(模型参数、优化器状态、训练进度等)
  • 主要用途:训练过程的中断与恢复
安全风险
  • 使用 pickle 序列化,存在代码注入风险
  • 缺乏内置的完整性验证机制
  • 难以检测潜在的恶意代码

3. .pt/.pth - PyTorch 原生格式

核心特性
  • 适用框架:PyTorch
  • 存储选项
    • 仅保存权重参数(state_dict)
    • 保存完整模型(包含结构)
  • 文件扩展名:.pt 和 .pth 本质相同,命名偏好不同
    • .pt 文件:更适合保存完整的模型,包括结构和参数。这使得在加载模型时不需要重新定义模型结构,便于模型的分享和部署。
      例如,当需要将模型直接用于推理时,.pt 文件是一个方便的选择。
    • .pth 文件:更适合保存模型的状态字典,通常用于模型的训练和开发阶段。
      当需要在不同环境中重新定义模型结构并加载参数时,.pth 文件更为灵活。
      例如,在迁移学习中,通常会加载 .pth 文件中的权重到新的模型结构中。
  • 文件大小
    • .pt 文件:通常较大,因为它可能包含完整的模型结构和参数。
    • .pth 文件:通常较小,因为它只包含模型的参数,不包含模型结构。
  • 社区偏好
    • .pt 文件:
      更常用于保存完整的模型,尤其是在模型部署和分享时。
    • .pth 文件:
      更常用于保存模型的状态字典,尤其是在模型训练和开发阶段。
  • 使用方法:
    # 保存模型权重
    torch.save(model.state_dict(), 'weights.pth')
    
    # 保存完整模型
    torch.save(model, 'model.pt')
    
    # 如果保存的是状态字典,需要先定义模型结构,再加载参数
    model = YourModelClass()
    model.load_state_dict(torch.load('weights.pth'))
    
    # 如果保存的是完整模型,可以直接加载模型权重
    model = torch.load('model.pt')
    

4. .onnx - 跨平台标准格式

核心特性
  • ONNX (Open Neural Network Exchange) 是一种开放的模型交换格式。
  • 设计目标:实现跨框架部署
  • 优化特性
    • 内置计算图优化
    • 广泛的硬件加速支持
    • 统一的模型表示
主要优势
  • 支持多种深度学习框架
  • 便于模型优化和部署
  • 良好的硬件适配性

5. .gguf - 高效本地推理格式

核心特性
  • 开发方:由 llama.cpp 项目引入,专为高效推理设计
  • 设计目标:优化本地推理性能
  • 特色功能:支持多种量化方案(如 4-bit、8-bit)
  • 优势
    • 高效推理:专为 CPU 和 GPU 混合推理设计。
    • 内容:单个文件包含模型的所有权重和配置信息。
    • 跨平台:适合在资源受限设备(如 Apple Silicon)上运行。
适用场景
  • 边缘设备部署
  • 资源受限环境
  • 需要量化优化的场景

6. .bin - 通用二进制格式

核心特性
  • 格式特点:非标准化的模型保存格式,但在某些情况下用于存储原始二进制权重数据
  • 使用灵活:可以灵活地组织权重文件,每个文件对应模型的不同层或组件
  • 主要用途:存储预训练的模型权重文件,这些文件包含了模型在训练过程中学到的参数值。通常是一系列的数字值,表示了神经网络中的连接权重、偏置项等参数
应用场景
  • 特殊部署环境
  • 自定义加载需求
  • 资源受限设备

7. .pkl - Python Pickle 序列化格式

核心特性
  • 实现基础:Python 的 Pickle 模块
  • 主要用途:Python 对象的序列化和反序列化
  • 应用范围:Python 生态系统
优势
  • 序列化操作简便直观
  • 支持复杂的 Python 对象结构
  • 与 Python 生态深度集成
安全风险
  • 反序列化攻击风险
    • Pickle 可执行任意代码
    • 易受恶意代码注入攻击
  • 最佳防护
    • 避免加载来源不明的 .pkl 文件
    • 考虑使用 .safetensors 等更安全的替代方案
    • 实施严格的文件来源验证

🔄 加载和保存方法

1. .safetensors 文件

  • 加载:使用 Hugging Face 提供的相关 API,如 safetensors.torch.load_file() 函数。
  • 保存:使用 Hugging Face 提供的保存函数,将模型的 state_dict 保存为 .safetensors 文件。

2. .ckpt 文件

  • 加载:在 PyTorch Lightning 框架中,通过定义一个继承自 pl.LightningModule 的类,并使用 pl.Trainer 来加载 .ckpt 文件。
  • 保存:在训练过程中,使用 trainer.save_checkpoint() 方法保存训练状态为 .ckpt 文件。

3. .pt/.pth 文件

  • 加载:使用 PyTorch 的 torch.load() 函数加载 .pth 文件,并通过 model.load_state_dict() 将加载的字典应用于模型实例。
  • 保存:使用 torch.save() 函数保存模型的 state_dict 或整个模型为 .pth 文件。

4. .onnx 文件

  • 加载:使用 ONNX 提供的 API 加载 .onnx 文件。
  • 保存:使用 ONNX 提供的 API 保存模型为 .onnx 文件。

5. .gguf 文件

  • 加载:使用 llama.cpp 提供的工具加载 .gguf 文件。
  • 保存:使用 llama.cpp 的转换工具将其他格式(如 .safetensors)转换为 .gguf 格式。

6. .bin 文件

  • 加载:需要自定义逻辑读取 .bin 文件,并将其中的权重应用到模型结构中。
  • 保存:同样需要自定义逻辑将模型的权重保存为 .bin 文件。

7. .pkl 文件

  • 加载:使用 Python 的 pickle 模块加载 .pkl 文件。
  • 保存:使用 Python 的 pickle 模块保存模型为 .pkl 文件。

🔧 格式转换指南

1. .safetensors.pth 之间的转换

  • 可以通过 Hugging Face 提供的 API 将 .safetensors 文件加载为 PyTorch 的 state_dict,然后使用 torch.save() 保存为 .pth 文件;反之亦然。

2. .ckpt.pth 之间的转换

  • 在 PyTorch Lightning 框架中,可以通过加载 .ckpt 文件恢复训练状态,然后提取出模型的 state_dict 并保存为 .pth 文件;或者将 .pth 文件中的 state_dict 加载到模型中,并使用 trainer.save_checkpoint() 保存为 .ckpt 文件。

3. .bin 文件的转换

  • 由于 .bin 文件是非标准化的格式,因此转换过程需要根据具体的 .bin 文件内容和结构来编写自定义代码。

4. .gguf 与其他格式的转换

  • 使用 llama.cpp 的转换工具将其他格式(如 .safetensors)转换为 .gguf 格式。

🚨 .ckpt 文件的安全风险与替代方案

尽管 .ckpt 文件在模型训练和部署中非常有用,但它存在以下安全隐患:

1. 恶意代码注入

  • .ckpt 文件使用 Python 的 pickle 模块进行序列化和反序列化。pickle 的一个显著缺点是,它允许在反序列化过程中执行任意代码。这意味着,恶意用户可以在 .ckpt 文件中嵌入恶意代码,当文件被加载时,这些代码会被执行,可能导致系统被入侵或数据泄露。
  • 例如,攻击者可以通过 pickle 注入恶意代码,加载 .ckpt 文件时,这些代码可能会窃取用户数据、植入木马病毒,甚至控制整个系统。

2. 缺乏安全性验证

  • 传统的 .ckpt 文件没有内置的安全机制来验证文件的完整性和来源。即使是从可信来源下载的文件,也可能在传输过程中被篡改,增加了安全风险。
  • 虽然一些平台(如 Hugging Face)提供了文件扫描功能,但这些工具并不能完全保证 .ckpt 文件的安全性。

3. 难以检测的恶意行为

  • 由于恶意代码可能隐藏在模型的权重数据中,常规的杀毒软件很难检测到这些威胁。即使检测到,也可能无法完全清除恶意代码。
  • 例如,某些恶意代码可能会在特定条件下触发,而常规扫描工具无法捕捉到这些隐蔽的行为。

🛡️ 更安全的替代方案:.safetensors

为了应对 .ckpt 文件的安全风险,Hugging Face 推出了 .safetensors 格式。以下是 .safetensors 的主要优势:

1. 安全性
  • .safetensors 文件仅包含模型的权重数据,不包含可执行代码。这种设计从根本上杜绝了恶意代码注入的可能性,确保模型加载过程的安全性。
  • .ckpt 文件不同,.safetensors 文件不支持 pickle 序列化,因此不会执行任何代码,大大降低了安全风险。
2. 高效性
  • .safetensors 文件采用优化的二进制格式,加载速度更快,内存占用更低。这对于处理大型模型(如数十亿参数的 GPT 模型)尤为重要。
  • 例如,在 Stable Diffusion 中,.safetensors 文件的加载速度比 .ckpt 文件快 30% 以上。
3. 跨平台兼容性
  • .safetensors 格式不依赖于特定的编程语言或框架,可以在多种深度学习工具中使用,提高了模型的通用性和可移植性。

📊 选择建议与格式对比

推荐使用场景

  1. 研发环境

    • 首选:.pt/.pth.safetensors
    • 备选: .ckpt(如果需要完整训练状态)
    • 原因:原生支持,调试方便
  2. 生产部署

    • 首选:.safetensors.onnx
    • 备选: .gguf(特别是边缘设备)
    • 原因:安全性高,性能优化
  3. 边缘设备

    • 首选:.gguf
  • 备选: .onnx
    • 原因:体积小,加载快
  1. 模型分享
    • 首选:.safetensors
    • 备选: .onnx、.pt/.pth
    • 原因:安全性高,兼容性好

性能指标对比

格式 加载速度 内存占用 安全性 跨平台性 量化支持 压缩率 版本兼容性
.safetensors ⭐⭐⭐⭐⭐ ⭐⭐⭐⭐ ⭐⭐⭐⭐⭐ ⭐⭐⭐⭐⭐ ⭐⭐⭐⭐ ⭐⭐⭐⭐
.pt/.pth ⭐⭐⭐ ⭐⭐⭐ ⭐⭐⭐ ⭐⭐⭐ ⭐⭐⭐ ⭐⭐⭐⭐
.ckpt ⭐⭐ ⭐⭐ ⭐⭐ ⭐⭐ ⭐⭐
.onnx ⭐⭐⭐ ⭐⭐⭐⭐ ⭐⭐⭐⭐ ⭐⭐⭐⭐⭐ ⭐⭐⭐ ⭐⭐⭐⭐⭐
.gguf ⭐⭐⭐⭐⭐ ⭐⭐⭐⭐⭐ ⭐⭐⭐⭐ ⭐⭐⭐⭐ ⭐⭐⭐⭐ ⭐⭐⭐
.bin ⭐⭐ ⭐⭐⭐ ⭐⭐ ⭐⭐ ⭐⭐ ⭐⭐
.pkl ⭐⭐ ⭐⭐⭐ ⭐⭐ ⭐⭐ ⭐⭐

详细特性对比

格式 主要应用生态 适用场景 核心优势 主要局限性 维护状态
.safetensors Hugging Face、PyTorch 生产部署、模型分享 1. 超高安全性
2. 优异加载性能
3. 完美跨平台支持
4. 内置序列化保护
1. 生态相对较新
2. 工具链不够丰富
活跃维护
.pt/.pth PyTorch 研发阶段、模型训练 1. PyTorch原生支持
2. 完整存储模型信息
3. 训练断点续传
1. 仅限PyTorch生态
2. 安全性一般
持续维护
.ckpt PyTorch Lightning 训练检查点、恢复训练 1. 完整训练状态保存
2. 支持训练恢复
3. 包含优化器状态
1. 安全性差
2. 文件体积大
3. 加载速度慢
逐渐减少
.onnx ONNX Runtime 跨平台部署、硬件加速 1. 广泛硬件支持
2. 跨框架兼容
3. 支持模型优化
1. 可能损失特性
2. 转换可能失败
3. 调试相对困难
活跃维护
.gguf llama.cpp 边缘设备、本地部署 1. 极致性能优化
2. 高效量化支持
3. 小内存占用
1. 主要面向推理
2. 不支持训练
3. 生态相对封闭
活跃维护
.bin 通用 自定义部署、资源受限 1. 格式简单
2. 灵活性高
3. 易于实现
1. 需要自定义加载
2. 缺乏标准化
3. 维护成本高
依项目而定
.pkl Python生态 快速原型、开发测试 1. Python原生支持
2. 序列化便捷
3. 使用简单
1. 安全性严重不足
2. 性能一般
3. 跨平台性差
保持稳定

🛡️ 安全性建议

  1. 文件来源验证

    • 仅使用可信来源的模型文件
    • 实施文件完整性校验
  2. 格式选择

    • 优先使用 .safetensors
    • 避免使用未经验证的 .ckpt 文件
  3. 部署安全

    • 实施访问控制
    • 定期安全审计

📈 未来趋势

  1. 安全性强化

    • .safetensors 格式将更加普及
    • 新的安全格式可能出现
  2. 效率提升

    • 更多针对特定硬件的优化格式
    • 混合精度存储方案
  3. 标准化

    • 统一的模型交换格式
    • 更好的跨平台兼容性

希望本文能帮助你更好地理解和选择适合的模型文件格式!如有问题,欢迎在评论区讨论。

Logo

技术共进,成长同行——讯飞AI开发者社区

更多推荐