AI 模型格式大揭秘

Java八股文面试

已于 2024-12-25 14:48:36 修改

阅读量1k

点赞数 27

文章标签： java 开发语言面试 spring 职场和发展

于 2024-12-25 14:24:18 首次发布

本文链接：https://blog.csdn.net/2401_89221867/article/details/144718876

版权

模型格式是什么

在机器学习中我们可以把模型想象成一个“黑盒”，我们给这个黑盒一些输入，然后它给我们一些输出，这个黑盒中会保存把输入转成输出的所有信息。这些信息一般来说我们需要固化持久化下来进行保存，以便后续使用。从这个角度来说模型是一种保存特定数据的文件文件，它承载着经过训练的模型的所有相关信息。它在整个机器学习（从训练到使用过程）中的作用十分关键，能够依据输入的数据，按照预先学习到的模式进行运算、推理，从而输出对应的结果，比如在图像识别领域可判断图像中的物体类别，在自然语言处理中能生成合适的文本回复等。

既然模型本质上只是一个或者多个文件，那它必然涉及到以什么样的组织形式来描述它以及存储它，事实上机器学习发展至今模型的存储经历了各种各样的格式，可谓是琳琅满目，从另一个方面揭示了机器学习的发展历程，有的是为了方便人类读写，有的是为了方便性能，有的的的确确是为了计算机存取方便。

比如：JSON格式、Pickle格式（转为 python）、PMML格式、HDF5格式（大量科学，复杂结构）、TensorFlow SavedModel格式、Torch Script格式、ONNX格式、PMML格式、Parquet格式（列式存储）。可见模型的格式有简单的文本格式也有二进制格式，有简单对象格式也有复杂结构格式。它们的诞生都有各自的作用，都在各自擅长的场景下发挥着过或者发挥过作用。

下面我讲介绍目前一些主流的格式。

ONNX

ONNX全称是Open Neural Network Exchange，由微软、FaceBook、亚马逊等多个公司一起推出的。它宛如一座桥梁，连接着不同的深度学习框架。它是一种开放的模型交换格式，为深度学习和传统机器学习模型定义了标准的文件格式以及运算符和数据类型。

ONNX采用了Protobuf（Google Protocol Buffer）格式进行存储，这种格式是是一种轻便高效的结构化数据存储格式，可以用于结构化数据序列化，很适合做数据存储或数据交换格式。

ONNX的主要特点

开放性：ONNX是一个开放标准，任何人都可以自由地使用它。
可扩展性：ONNX支持多种数据类型和操作，可以表示复杂的神经网络模型。
互操作性：ONNX可以在不同的深度学习框架和平台之间无缝迁移模型。
性能：ONNX旨在提供高效的模型表示，以便在不同的硬件上实现快速的推理。
优势：
- 模型共享：使得不同框架之间的模型共享成为可能。
- 部署便利：简化了模型在不同环境中的部署过程。
- 优化支持：ONNX支持多种优化技术，以提高模型在推理时的性能。

从工作原理来看，ONNX 的核心是中间表示（IR）格式，它清晰地定义了一组标准运算符和数据类型，这种 IR 格式能够被各种深度学习框架和工具轻松理解与处理。当开发者使用支持 ONNX 的框架训练模型后，可将模型导出为 ONNX 格式，随后，该 ONNX 模型便能被其他支持 ONNX 的框架或运行时加载并执行。

一个 ONNX 模型通常包含几个关键组件：

图（Graph）精准地描述了模型的整体结构，涵盖输入、输出以及中间计算节点；
节点（Node）代表着模型中的各类操作，如卷积、池化等；
张量（Tensor）用于存储和传递数据的多维数组；
属性（Attribute）则明确了节点的特定参数。

如此结构设计，使得 ONNX 模型具备良好的可读性与可移植性。

ONNX 的生态系统也十分丰富，包括 ONNX 格式、ONNX Runtime（一个高性能的推理引擎，可在多种平台上运行 ONNX 模型）、ONNX 工具（如模型转换、优化和可视化等工具）以及社区贡献（来自全球开发者的各种工具、库和扩展）。这一生态系统的存在，让 ONNX 不仅仅是一种文件格式，更成为一个完整的 AI 模型交换和部署解决方案，在模型部署、硬件优化、模型共享以及生产环境适配等实际应用场景中展现出巨大价值。

TensorFlow 模型格式

TensorFlow 框架下可导出的模型格式，主要包括两种：SavedModel 和 HDF5。其中SavedModelSavedModel 是 TensorFlow 推荐的序列化格式，用于保存和加载训练好的模型，里面不仅装着模型的参数权值，还完整地囊括了模型的计算流程。这使得它在模型部署和跨平台使用时表现得极为出色，备受开发者们的青睐。它支持多种语言，包括 Python、C++ 和 Java 等。SavedModel 包含了模型的结构、权重数据和优化器状态等信息，可以在不同的环境中进行部署和使用

SavedModel 结构

SavedModel 的主要结构包括以下几个文件：

saved_model.pb：模型的 GraphDef 文件，包含了模型的计算图结构。
variables/：包含模型权重的变量文件。
assets/：包含模型的额外资源文件，如文本和图像等。
meta_graph.pbtxt：包含模型的元数据文件，如签名函数和输入输出信息等。

SavedModel 使用

当我们使用 TensorFlow 训练好一个模型后，要将其保存为 SavedModel 格式，其实并不复杂。在 TensorFlow 1.x 版本中，我们可以借助 tf.saved_model.builder.SavedModelBuilder 来精心构建一个 SavedModel。比如，首先创建一个 SavedModelBuilder 对象，指定保存的路径，接着定义好模型的输入和输出签名（Signature），这一步就像是给模型的入口和出口贴上清晰的标签，方便后续使用。然后，将计算图（Graph）和变量（Variables）添加进去，最后使用 save 方法就能将模型完整地保存下来。而在 TensorFlow 2.x 版本中，使用 tf.saved_model.save 函数就可以更加简洁地完成保存操作，这就像是按下了一个快捷保存按钮，轻松搞定模型保存。

加载 SavedModel 同样便捷，在 Python 中，使用 tf.saved_model.load 函数就能将模型加载进来，并且可以直接使用模型进行推理预测，就像打开宝盒就能直接取用里面的宝贝一样方便。例如，在一个图像分类任务中，我们加载已经保存好的图像分类模型，然后输入一张新的图片，模型就能快速给出这张图片的分类结果。

python

代码解读

复制代码

import tensorflow as tf model = tf.keras.models.load_model('path/to/saved_model')

与 TensorFlow 其他的模型保存格式相比，SavedModel 具有显著的优势。和 checkpoint 格式相比，checkpoint 主要用于保存模型训练过程中的中间参数值，它就像是模型训练过程中的一个个快照，方便在训练中断后继续训练，但它并不包含完整的计算图信息，这就使得在模型部署和跨平台使用时不够方便。而 SavedModel 则完整地保存了模型的所有信息，无论是进行推理还是部署到不同的平台上，都更加得心应手。与基于protobuf 格式保存的模型相比，SavedModel 在易用性和兼容性方面表现得更为出色。protobuf 虽然也是一种常用的数据序列化格式，但在实际使用中，SavedModel 能够更好地与 TensorFlow 生态系统中的其他工具和库进行无缝对接，为开发者提供更加流畅的开发体验。

PyTorch.pth

在 PyTorch 的世界里，主要以.pth 文件保存模型，可以用来保存模型的权重，让模型的权重得以存储。.pth文件格式是一种用于存储深度学习模型参数的文件格式。这种格式设计简洁且高效，便于模型的保存、加载和迁移学习。我们可以使用 torch.save 函数，把模型的学习成果保存起来，如果传入的是模型的 state_dict（模型参数的字典），那就是只保存了模型的参数信息；要是直接传入模型本身，那就连模型的结构一起保存下来了。

文件结构

.pth文件通常包含一个或多个数据块（blobs），每个数据块存储了模型中特定层的参数。数据块的名称通常与模型中的层名称相对应。例如，一个简单的卷积神经网络可能包含以下数据块：

conv1.weight：卷积层1的权重
conv1.bias：卷积层1的偏置
conv2.weight：卷积层2的权重
fc1.weight：全连接层1的权重
fc1.bias：全连接层1的偏置

举个例子，我们训练了一个简单的卷积神经网络用于图像分类，使用 torch.save(model.state_dict(), 'my_model.pth') 就把模型学习到的卷积核权重、全连接层权重等参数保存成了 .pth 文件。而加载 .pth 文件也不难，对于只保存了参数的 .pth 文件，我们得先创建一个相同结构的模型实例，就像先准备一个空的书架，然后用 torch.load 函数把文件中的参数加载到这个模型实例中，使用 model.load_state_dict(torch.load('my_model.pth')) 语句，模型就像被注入了灵魂，能够进行推理预测了。

python

代码解读

复制代码

import torch import torchvision.models as models # 加载预训练模型 model = models.resnet18(pretrained=True) # 保存模型到.pth文件 torch.save(model.state_dict(), 'model.pth')

python

代码解读

复制代码

# 加载保存的模型参数 model = models.resnet18() model.load_state_dict(torch.load('model.pth')) # 设置模型为评估模式 model.eval()

.pth 文件在实际应用中用途广泛。在模型部署方面，我们可以把训练好的模型保存为 .pth 文件，然后在不同的环境中加载它，让模型在生产环境中发挥作用，比如将图像分类模型部署到手机 APP 中，实现对用户拍摄照片的实时分类。在迁移学习中，我们可以加载预训练模型的 .pth 文件，然后在其基础上进行微调，就像是站在巨人的肩膀上继续探索，比如使用在大规模图像数据集上预训练的模型，针对特定的医学图像数据集进行微调，从而快速适应新的任务。而且，研究人员之间也经常共享 .pth 文件，这样大家就可以方便地复现和改进模型，促进学术交流与合作。

不过，使用 .pth 文件也有一些需要注意的地方。版本兼容性是个重要问题，如果 PyTorch 版本更新了，可能会导致旧版本保存的 .pth 文件在加载时出现不兼容的情况，就像一把旧钥匙可能打不开新锁。所以在保存模型时，最好记录下使用的 PyTorch 版本信息，以便后续排查问题。此外，文件安全性也不容忽视，因为 .pth 文件可能包含敏感信息，如模型的权重参数，如果被恶意获取，可能会导致模型被窃取或篡改，所以在保存和传输 .pth 文件时，要采取适当的加密和安全措施。而且，随着模型复杂度的增加，.pth 文件的大小也可能会变得很大，这就需要我们考虑如何优化存储和传输，比如采用模型压缩技术，减小文件大小，提高传输效率。

另外，在保存和加载模型时，需要确保模型的架构（如层名称和顺序）与保存时一致，否则可能会导致加载失败或参数不匹配的问题。对于复杂的模型，.pth文件可能会变得非常大，因此在实际应用中可能需要考虑使用更高效的存储格式，如ONNX。

safetensors 格式

Safetensors是一种由Hugging Face团队开发的新格式，用于安全地存储张量（tensors），它是PyTorch的pickle格式的替代品，因为pickle格式不安全，可能包含可以执行的恶意代码。Safetensors格式的文件通常以.safetensors为扩展名，它通过避免序列化任意代码的方式，确保在加载文件时不会执行恶意代码，从而提供了更高的安全性。