大模型参数文件存储格式全解析

作者:十万个为什么2024.11.21 13:58浏览量:195

简介:文章深入解析了大模型中.safetensors、.ckpt、.pth和.bin文件的区别、加载保存方法以及转换方式,帮助用户更好地理解和应用这些文件格式。

深度学习领域,模型的参数和权重通常以各种文件格式进行存储,以便在训练、推理和部署过程中快速加载和使用。其中,.safetensors、.ckpt、.pth和.bin文件是四种常见的存储格式。本文将详细解析这四种文件的区别、加载保存方法以及转换方式。

一、文件格式区别

  1. .safetensors文件

    • 来源:由Hugging Face团队开发。
    • 用途:专为存储和加载大型张量(tensors)而设计,特别关注模型安全性、隐私保护和快速加载。
    • 内容:仅包含模型的权重参数,不包括执行代码。
    • 优势:通过优化的二进制格式和支持内存映射,显著提高大模型文件的加载速度;文件布局设计使得读取特定张量时无需加载整个文件;提供丰富的API和工具,便于与主流深度学习框架集成。
  2. .ckpt文件

    • 来源PyTorch Lightning框架采用的模型存储格式。
    • 用途:用于保存和恢复训练过程,包括模型参数、优化器状态以及可能的训练元数据信息。
    • 内容:包含完整的训练状态,使得用户可以无缝地恢复训练或执行推理。
    • 优势:方便存储和恢复训练状态,节省时间;可以记录模型的不同版本,便于比较和选择。
  3. .pth文件

    • 来源:PyTorch中用于保存模型状态的标准格式。
    • 用途:保存模型的state_dict(包含所有可学习参数)或整个模型(包括结构和参数)。
    • 内容:可以是模型的权重参数或整个模型实例。
    • 优势:PyTorch的通用模型保存和加载格式,跨版本兼容性好;可以方便地分享和部署PyTorch模型。
  4. .bin文件

    • 来源:非标准化的模型保存格式,但在某些情况下用于存储原始二进制权重数据。
    • 用途:存储预训练的模型权重文件,这些文件包含了模型在训练过程中学到的参数值。
    • 内容:通常是一系列的数字值,表示了神经网络中的连接权重、偏置项等参数。
    • 优势:可以灵活地组织权重文件,每个文件对应模型的不同层或组件。

二、加载和保存方法

  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. .pth文件

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

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

三、转换方式

由于.safetensors、.ckpt、.pth和.bin文件在格式和内容上存在差异,因此它们之间的转换通常需要借助一些工具或编写自定义代码。

  • .safetensors与.pth之间的转换:可以通过Hugging Face提供的API将.safetensors文件加载为PyTorch的state_dict,然后使用torch.save()保存为.pth文件;反之亦然。
  • .ckpt与.pth之间的转换:在PyTorch Lightning框架中,可以通过加载.ckpt文件恢复训练状态,然后提取出模型的state_dict并保存为.pth文件;或者将.pth文件中的state_dict加载到模型中,并使用trainer.save_checkpoint()保存为.ckpt文件。
  • 与.bin文件的转换:由于.bin文件是非标准化的格式,因此转换过程需要根据具体的.bin文件内容和结构来编写自定义代码。

四、产品关联

在深度学习模型的开发和部署过程中,选择一个合适的平台或工具至关重要。千帆大模型开发与服务平台提供了全面的模型开发、训练和部署功能,支持多种文件格式(包括.safetensors、.ckpt、.pth和.bin)的加载和保存。通过使用千帆大模型开发与服务平台,用户可以更加高效地管理模型参数和权重文件,加速模型的训练和推理过程。

例如,在千帆大模型开发与服务平台上,用户可以轻松地将训练好的模型保存为.safetensors文件,以便在需要时快速加载和部署。同时,平台还支持将.ckpt或.pth文件转换为其他格式,以满足不同场景下的需求。此外,平台还提供了丰富的API和工具,帮助用户更好地管理和优化模型参数和权重文件。

综上所述,.safetensors、.ckpt、.pth和.bin文件在深度学习领域各有其独特的用途和优势。通过了解它们的区别、加载保存方法以及转换方式,用户可以更加灵活地管理模型参数和权重文件,提高模型的训练和推理效率。