PyTorch 与 TensorFlow 模型保存格式全解析与代码示例

前言

在深度学习模型训练完成后，如何保存和加载模型是每个开发者都会面临的问题。不同框架提供了多种模型保存格式，每种格式都有其特点和适用场景。本文将面向中级开发者，系统介绍 PyTorch 和 TensorFlow 中常见的模型保存格式，包括 PyTorch 的 .pth/.pt（原生格式）和 ONNX，以及 TensorFlow 的 SavedModel、HDF5 (.h5)、协议缓冲文件 (.pb) 和 TensorFlow Lite (.tflite)。我们将解析每种格式所包含的内容，给出保存与加载的代码示例，比较它们的区别，讨论转换方法，并分享常见坑点与最佳实践建议。

PyTorch 模型保存格式综述

PyTorch 通常使用二进制文件来保存模型参数或整个模型。常见扩展名是 .pt 或 .pth，二者在本质上没有区别，都是通过 PyTorch 的序列化机制（pickle 或 zip格式）保存数据。此外，PyTorch 还支持将模型导出为跨平台的 ONNX 格式，以便在其他环境中运行。下面分别介绍这些格式。

PyTorch 原生格式（.pth / .pt）

1. 特点与适用场景： PyTorch 原生格式使用自有的序列化方式保存模型。通常开发者有两种保存策略：只保存模型参数（state dict）或保存整个模型结构+参数。官方推荐使用保存 state dict 的方式，因为这种方式灵活且与代码解耦​。.pth 和 .pt 仅是习惯用法上的扩展名，均可使用；其中 .pt 常被使用于保存经过 TorchScript 序列化的模型（便于C++部署），而 .pth 常用于保存纯 Python 下的模型权重文件。一般来说，在 PyTorch 中保存模型用于后续继续训练或在相同代码下推理时，使用 .pth（状态字典）最佳；而在部署推理场景（例如C++环境）下，可以将模型转换为 TorchScript 并保存为 .pt。
2. 保存内容解析： 取决于保存方式，.pth/.pt 文件可以包含不同信息：

• 仅参数（state dict）：文件仅包含模型的权重张量映射（键为参数名，值为参数张量）。不包含模型架构或定义，需要在加载时先定义模型类，再加载参数。​如果一起保存了优化器的 state dict，则还包含优化器的动量等状态。
• 整个模型（结构+参数）：通过直接 torch.save(model, ... ) 保存时，文件包含模型对象完整序列化（包括模型架构和参数）。PyTorch 使用 Python 的 pickle 序列化模型对象​。这种方式使文件与代码强绑定 ，即加载时依赖于原始模型类的定义，若类定义改变或路径不同，将无法正常加载​。因此不利于跨项目使用或后续代码重构。

PyTorch 保存 & 加载示例

通常推荐保存和加载 state dict：

import torch
import torch.nn as nn

model = ...  # 假设已定义并训练的模型

# 保存模型参数（state_dict）
torch.save(model.state_dict(), "model_weights.pth")

# ------------ 稍后在同环境下加载 -------------
model = ...  # 先重新定义模型结构并实例化
model.load_state_dict(torch.load("model_weights.pth"))
model.eval()  # 切换为推理模式

上面代码将模型参数保存为 model_weights.pth。加载时需要先初始化相同结构的模型，然后使用 load_state_dict 加载参数并调用 eval() 切换为评估模式（确保如 Dropout、BatchNorm 层正确工作）。

如果希望保存完整模型（不想手动重建模型结构），可以这样：

# 保存整个模型
torch.save(model, "model_full.pth")

# ------------ 加载整个模型 -------------
model_full = torch.load("model_full.pth")
model_full.eval()

这会直接保存模型对象到文件，加载时无需手动定义模型类。但是注意，此方法依赖于保存时的模型类定义：
注意: 使用 torch.save(model) 保存完整模型时，PyTorch 会通过 pickle 序列化模型对象。这要求在加载时存在完全相同的模型类定义和目录结构。如果代码发生改变或在另一个项目中加载，该方法可能失败​。因此除非特殊需要，一般更推荐仅保存 state dict 并在代码中维护模型结构。

TorchScript 序列化 (.pt)：PyTorch 提供 TorchScript 来将模型转换为可独立运行的序列化格式。通过 torch.jit.trace 或 torch.jit.script，可以将训练好的模型转换为 TorchScript Module，然后保存为 .pt 文件：

# 将模型转换为 TorchScript 并保存
scripted_model = torch.jit.trace(model, example_inputs=torch.randn(1, 3, 224, 224))
scripted_model.save("model_scripted.pt")

生成的 model_scripted.pt 文件包含模型的结构和权重，但采用 TorchScript 表示，其优点是可以在纯 C++ 环境下加载并运行（通过 LibTorch），无需 Python 解释器。这种格式适合部署到服务端或移动端，高性能推理环境，是 PyTorch 官方推荐的部署模型格式​。TorchScript 模型同样不包含优化器等训练状态，仅用于推理。

PyTorch ONNX 格式

1. 特点与适用场景： ONNX（Open Neural Network
   Exchange）是一种开放的神经网络交换格式，旨在让模型在不同深度学习框架之间转换和共享​。通过 ONNX，我们可以将PyTorch 训练好的模型导出为一个标准格式的文件（扩展名 .onnx），然后在其他支持 ONNX 的推理引擎或框架中运行，例如ONNX Runtime、TensorRT，或者转换为 TensorFlow 等。ONNX 非常适用于跨平台和跨框架部署，例如需要将 PyTorch 模型在非 Python环境（C++、Java）中推理，或者在移动设备上通过 Windows ML、OpenCV DNN 等调用。
2. 保存内容解析： ONNX 文件使用protobuf（二进制协议缓冲）格式存储，内部包含模型的计算图定义（网络的算子及连接）和训练后的权重值。与 PyTorch 原生格式不同，ONNX 不包含任何 Python 代码或类信息，也不保存优化器状态等训练元数据，仅关注推理所需的网络结构和参数。这意味着：

• ONNX 不依赖原框架代码即可运行模型（只需合适的推理引擎即可）。
• 由于只保存推理相关信息，它无法直接用于在原框架中继续训练或恢复优化器状态。
• PyTorch 当前无法直接导入 ONNX文件重新得到模型（需要通过其他方式加载），因此 ONNX 通常是单向导出用于部署的格式​。

PyTorch 导出 ONNX 示例：

import torch
import torchvision.models as models

model = models.resnet18(pretrained=True)
dummy_input = torch.randn(1, 3, 224, 224)  # 示例输入张量

# 导出模型为 ONNX 文件
torch.onnx.export(model, dummy_input, "resnet18.onnx", 
                  input_names=["input"], output_names=["output"])
print("ONNX 模型已保存为 resnet18.onnx")

以上代码使用 torch.onnx.export 将模型导出为 ONNX。需要提供一个示例输入 (dummy_input) 来跟踪模型的计算图。这会在当前目录生成 resnet18.onnx 文件。

ONNX 模型加载与使用示例： 虽然 PyTorch 不能直接加载 ONNX 文件执行，但我们可以借助 ONNX Runtime 或其他引擎：

import onnx
import onnxruntime as ort
import numpy as np

# 加载 ONNX 模型并检查
onnx_model = onnx.load("resnet18.onnx")
onnx.checker.check_model(onnx_model)  # 验证模型结构正确性

# 使用 ONNX Runtime 执行推理
ort_sess = ort.InferenceSession("resnet18.onnx", providers=['CPUExecutionProvider'])
# 准备 numpy 输入
x = np.random.randn(1, 3, 224, 224).astype(np.float32)
# 获取 ONNX 模型的输入名（这里我们导出时命名为 "input"）
input_name = ort_sess.get_inputs()[0].name
# 推理
result = ort_sess.run(None, {input_name: x})

上述代码演示了如何用 ONNX Runtime 加载 .onnx 模型并进行推理。通过 ONNX，我们可以在没有 PyTorch 环境的情况下运行模型。但需要注意不同推理引擎对算子的支持情况，导出时应选择合适的 opset 版本，并在导出后验证 ONNX 模型输出是否与原始 PyTorch 模型一致（某些算子可能在不同平台表现略有差异，需仔细测试）。

TensorFlow 模型保存格式综述

TensorFlow 提供了多种格式来保存模型，主要有 TensorFlow 2.x 默认的 SavedModel 格式、Keras API 常用的 HDF5 (.h5) 格式，以及 TensorFlow 1.x 时代常见的 GraphDef .pb 文件（冻结模型）格式。对于移动端和嵌入式部署，还有 TensorFlow Lite (.tflite) 格式。下面依次介绍各格式的特点和使用方法。

TensorFlow SavedModel 格式

1. 特点与适用场景： SavedModel 是 TensorFlow 自 1.x 后期推出并在 2.x成为默认的模型保存格式。它能够保存完整的 TensorFlow程序，包含训练好的参数（即变量的值）和计算图等所有内容。SavedModel不依赖原始构建模型的代码即可恢复模型的计算过程​。这使得它非常适合分享和部署：例如使用 TensorFlow Serving部署到服务器、在 TensorFlow.js 或 TensorFlow Lite 中加载，或发布到 TensorFlow Hub。SavedModel 保存后是一个目录，其中包含了模型的元数据和数据文件（如 saved_model.pb 和 variables/ 子目录等）。
2. 保存内容解析： SavedModel 包含以下主要内容：

• 计算图 / 函数：以协议缓冲格式存储在 saved_model.pb 文件中，对应模型的计算逻辑（在 TF2.x 中为通过 tf.function 转换的计算图）。SavedModel 实际上保存了模型的执行图（而非简单的配置)，因此连自定义的模型结构或算子也能被保存下来，无需原始 Python 定义即可运行​。
• 权重参数：保存在 variables/ 目录下的二进制文件（如 variables.data-00000-of-00001 和 variables.index），表示所有 tf.Variable 的取值。
• 签名(signature)：模型的输入输出接口定义（函数签名），便于在部署时调用正确的方法。例如默认推理函数叫 serving_default，列出了输入输出张量的名称和形状。
• 额外资产：如果模型需要额外的文件（如词汇表、label映射等），可放在 assets/ 目录中，由 SavedModel 一并打包。

SavedModel 是一种语言无关、跨平台的格式，任何支持 TensorFlow 的环境都可以加载并执行 SavedModel 中的图。由于保存的是计算图本身，它相比其它格式更健壮完整，但保存的体积可能稍大，同时可读性较低（内容是序列化的GraphDef）。SavedModel 特别适合需要长期保存或跨环境部署的情况，例如模型最终上线部署、通过 C++ TF Serving 提供服务、转换为 TensorFlow Lite 或 TensorFlow.js 等。

保存与加载示例（Keras API）

import tensorflow as tf
from tensorflow import keras

model = keras.Sequential([...])  # 构建或获取一个模型
# 训练模型...
# 保存模型为 SavedModel 格式（保存为目录）
model.save("saved_model_dir")  # 默认不指定格式将保存为SavedModel

上面的代码会在当前目录生成一个名为 saved_model_dir 的文件夹。其中包含上述提到的 saved_model.pb、variables 等子文件。加载SavedModel同样可以使用Keras API：

# 加载 SavedModel
model = keras.models.load_model("saved_model_dir")

tf.keras.models.load_model 可以自动识别目录中的 SavedModel 并加载为 Keras 模型对象（前提是该 SavedModel 是通过 Keras 保存的或包含Keras网络结构）。如果 SavedModel 不是来自 Keras模型（例如自己用低级API构建并tf.saved_model.save保存），也可以使用低级 API 加载：

loaded = tf.saved_model.load("saved_model_dir")
infer = loaded.signatures["serving_default"]
result = infer(tf.constant(input_data))

上述方式直接加载为一个 Trackable 对象，用 signatures 获取推理函数来执行。总的来说，SavedModel 提供了高低级两套加载接口：对 Keras模型友好的 keras.models.load_model 返回原Keras模型（可继续训练等），以及底层的 tf.saved_model.load 返回可用于推理的函数集合。

SavedModel 文件结构示例： （保存后目录可能如下）

这种格式在TensorFlow 2.x中被官方广泛使用和支持，是默认和推荐的模型保存格式，特别当你需要在训练完成后部署模型或与他人分享模型时。

Keras HDF5 (.h5) 格式

1. 特点与适用场景： HDF5 是一种通用的分层数据格式。在 TensorFlow (尤其是 Keras API) 中，使用 .h5 扩展名可以将模型保存为 HDF5 文件。这个格式在 TensorFlow 1.x 和早期 Keras 中非常常见，也在 TensorFlow 2.x 中继续受到支持。与 SavedModel 相比，HDF5 以单文件形式保存模型，使用直观，文件便于传输。适合快速保存/加载模型权重，或者需要与不支持SavedModel的旧工具兼容的场景​。许多只关注模型架构和权重的小型项目会选择保存 .h5 文件以便管理。
2. 保存内容解析： 一个通过 Keras 保存的 .h5 文件通常包含：

• 模型架构：如果模型是通过Sequential或函数式API构建，结构会以JSON的形式存储在HDF5中，使得加载时能够重建模型拓扑。
• 权重参数：模型各层的参数以数据集形式存储在HDF5文件内部。
• 训练配置：如果模型在保存前经过编译（model.compile()），优化器的名称、学习率等配置会被保存。
• 优化器状态：如果模型经过了一定训练，保存时包含优化器的内部状态（比如动量项、二阶矩等），这样加载后可以从中断处继续训练​。

换言之，HDF5 格式可以看作将模型的配置+参数+训练状态打包到了一个文件中​。但是需要注意，HDF5 保存模型架构是基于模型的配置对象（config）。对于使用 Keras 函数式或序贯模型可以完美记录结构。但对于自定义的Subclassing API（继承 Model 自定义模型）而言，架构很难序列化为配置，这种情况下调用 model.save('x.h5') 可能无法保存模型结构。TensorFlow 会在保存时给出警告，并退而仅保存权重（或在较新版本中自动切换SavedModel）。所以在复杂自定义模型场景，SavedModel 更可靠​。

保存 & 加载示例

# 继续使用前面的 model（Keras模型）
model.save("model.h5")  # 保存为单个 HDF5 文件
# ... 之后 ...
model = keras.models.load_model("model.h5")

一行 model.save("model.h5") 就能将模型保存为 HDF5 文件。稍后用 load_model 即可加载回来。这对于标准的 Sequential 或 Functional 模型非常方便，加载后模型已经编译好且包含之前的权重和优化器状态，可以直接 .compile() 后 fit 继续训练。

HDF5 与 SavedModel 异同： 总的来说，在 TensorFlow 2.x 中，SavedModel 是官方默认格式，但 HDF5 依然可用：

• 文件数量：HDF5 是单文件，SavedModel 是目录（多个文件）。
• 可移植性：SavedModel 保存执行图，可跨语言跨平台直接运行​；HDF5 需要在支持 Keras 的环境解析模型配置，某种程度上依赖代码（尤其对自定义对象）。
• 信息完整度：二者都可保存权重和训练配置，但 SavedModel 能涵盖自定义计算图（例如tf.function，TensorFlow Hub等整个Graph），HDF5 更关注模型结构本身​。
• 推荐场景：SavedModel 用于部署、服务化和需要充分跨环境复现的情况；HDF5 常用于本地开发调试或简单模型的保存，或者需要与旧版Keras(例如TensorFlow 1.x, Theano后端的Keras)兼容的情况​。实际使用中，TensorFlow 2.x 如果不特别指定，一般会保存SavedModel；只有在明确要求单文件时才使用.h5。

冻结模型 GraphDef (.pb) 格式

1. 特点与适用场景： .pb 是 TensorFlow 协议缓冲(Protobuf) 格式模型文件的扩展名，通常用来表示计算图定义（GraphDef）。在 TensorFlow 1.x 中，模型的计算图和权重存储分离：计算图可以导出为一个 .pb 文件（GraphDef)，而权重保存在 checkpoint (.ckpt) 文件中。如果需要将模型用于推理部署，通常会将 GraphDef 和权重合并，生成一个冻结的模型文件（freeze graph），其本质仍是 .pb 格式，但图中的变量已转换为常量并嵌入权重值。这样 .pb 文件即可独立用于推理，不再依赖 checkpoint。

在 TensorFlow 2.x 中，冻结图的概念被 SavedModel 所取代（SavedModel 内部也有 .pb 文件），但在某些场景下开发者仍会接触 .pb：

• 使用一些第三方推理框架或工具需要导入 .pb（例如 OpenCV DNN、一些老的移动端API等）。
• 从 TensorFlow 2.x 模型获取冻结图用于兼容旧代码或作进一步转换（如转换为 CoreML、NCNN 等格式）。

2. 保存内容解析： .pb 文件（GraphDef 或冻结模型）包含模型的计算图结构，以及在冻结模型情况下包含每层的参数值作为 GraphDef 中 Const 节点的一部分。它不包含优化器状态或任何关于训练的信息，也不直接保存模型的高级配置。可以认为 .pb 主要服务于推理用途，一旦模型保存为冻结图，就无法再继续训练（因为变量已变为常量）。同时，由于 GraphDef 缺乏高级结构信息，不能通过它重建 Keras 模型来调用 .fit 等——如果需要进一步训练，一般应保留 SavedModel 或 HDF5 格式。
3. 冻结模型导出示例（TF2.x获取冻结GraphDef）: 在 TF2 中，导出 .pb 不是直接调用 model.save 能得到的，需要以下步骤：

• 将模型保存为 SavedModel（得到包含 saved_model.pb）。
• 使用 tf.saved_model.load 加载模型并获取 ConcreteFunction（具体的计算图函数）。
• 将该函数转换为 GraphDef（通过 function.graph.as_graph_def() 等方法），并使用 tf.io.write_graph 保存为 .pb 文件。

简要代码示例

import tensorflow as tf
model = ... # 构建并训练模型
tf.saved_model.save(model, "tmp_savedmodel")
# 加载刚保存的模型
loaded = tf.saved_model.load("tmp_savedmodel")
concrete_func = loaded.signatures["serving_default"]
# 冻结图：将变量转换为常量
frozen_func = tf.graph_util.convert_variables_to_constants_v2(concrete_func)
graph_def = frozen_func.graph.as_graph_def()
# 保存 GraphDef 为 .pb 文件
tf.io.write_graph(graph_def, ".", "frozen_model.pb", as_text=False)

执行后会在当前目录生成一个 frozen_model.pb。这个文件就是独立的冻结模型，可以在需要 .pb 的环境下使用（例如使用 OpenCV 的 cv2.dnn.readNetFromTensorflow("frozen_model.pb") 加载等）。 需要强调，SavedModel 已经包含了 .pb（GraphDef）和变量，本质上比冻结 .pb 更灵活。因此除非有特定需求（比如兼容旧库），大多数情况下优先使用 SavedModel 而非手动冻结图。

TensorFlow Lite (.tflite) 格式

1. 特点与适用场景： TensorFlow Lite (TFLite) 是针对移动端、嵌入式设备优化的模型格式和运行时。TFLite 使用特殊的FlatBuffer 二进制格式存储模型，文件扩展名为 .tflite​。TFLite 模型体积小、加载快，专为设备上高效执行而设计。例如，在手机、物联网设备上运行离线推理，用 TFLite 格式可以降低延迟和内存占用​。TensorFlow 提供了转换器将现有 TensorFlow 模型（SavedModel、Keras模型或冻结图）转换为 .tflite 文件，并配套提供 TFLite 解释器来运行这些模型。
2. 保存内容解析： .tflite 文件内部采用 FlatBuffer 表示，包含了模型的算子逻辑和权重参数，类似精简的计算图。相较于标准 TensorFlow 的 Protocol Buffer 格式，FlatBuffer 不需要解析/反序列化即可随机访问，因而加载和推理速度更快，而且模型文件通常更小​。TFLite 模型同样只面向推理阶段，不包含训练相关内容（优化器、gradients等均无）。另外，TFLite 格式支持量化和裁剪等优化：转换过程中可以将浮点模型量化为 8-bit 整数等，以进一步减小模型体积和提高推理效率。

转换与加载示例

将 Keras 模型或 SavedModel 转换为 .tflite：

import tensorflow as tf

model = ...  # 假设已经训练好的 tf.keras 模型
# 使用 TFLiteConverter 从 Keras 模型创建转换器
converter = tf.lite.TFLiteConverter.from_keras_model(model)
# （可选）转换优化，如开启量化：
converter.optimizations = [tf.lite.Optimize.DEFAULT]
# 进行转换
tflite_model = converter.convert()
# 将结果写入文件
with open("model.tflite", "wb") as f:
    f.write(tflite_model)
print("TFLite 模型已保存为 model.tflite")

TFLite 没有像 tf.keras 那样直接 .load_model 的接口，而是使用专门的解释器运行模型：

import numpy as np
import tensorflow.lite as tflite

# 加载 TFLite 模型
interpreter = tflite.Interpreter(model_path="model.tflite")
interpreter.allocate_tensors()
# 获取模型输入输出的详细信息
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
# 准备输入数据并喂入模型
input_data = np.array(..., dtype=np.float32)
interpreter.set_tensor(input_details[0]['index'], input_data)
interpreter.invoke()  # 执行推理
# 提取输出结果
output_data = interpreter.get_tensor(output_details[0]['index'])

上述代码通过 Interpreter 加载 .tflite，然后设置张量并调用 invoke() 运行推理。TFLite 模型非常小巧且独立，可在Android/iOS 应用、微控制器等环境中由对应平台的API加载执行。它支持多种语言（Java、C++、Python等）调用推理。

值得注意的是，.tflite 通常结合硬件加速（如利用移动端的NNAPI、GPU委托等）以获得最佳性能​。因此如果目标是移动部署或边缘设备，使用 TFLite 格式是最佳选择。

各格式对比与转换总结

以上已经介绍了 PyTorch 与 TensorFlow 框架下多种模型保存格式，它们各有优劣。下面对主要格式做一个对比总结，并给出格式选择和转换的建议：

格式选择建议

开发过程中，可以综合考虑需求来选择格式：

PyTorch 与 TensorFlow 之间直接转换并不简单，需要借助 ONNX 作为桥梁（例如 PyTorch -> ONNX -> TensorFlow，或 TensorFlow -> ONNX via tf2onnx）。而 TensorFlow 内部的 SavedModel 和 HDF5 转换相对容易：TF2.x 中通过 model.save() 参数 save_format='h5' 可以直接保存 HDF5；反之加载后再 model.save(save_format='tf') 可得 SavedModel。GraphDef .pb 则可从 SavedModel 得到或转换为 SavedModel，但涉及底层操作。

总之，没有一种格式适合所有场景。开发者应根据需求选择：在训练阶段注重可继续训练和易调试，在部署阶段强调自包含和跨平台。理解各格式的内容和差异，有助于做出正确选择并进行必要的格式转换。

实际操作建议

模型保存不是训练结束才考虑的事情，从项目一开始就应规划模型保存策略。例如，决定是每epoch保存checkpoint、还是只保存最佳模型；用何种格式保存以兼顾调试和部署需求。对于重要的里程碑模型，建议同时保存 两种格式：在TensorFlow中可以同时保存 SavedModel 和 .h5，在PyTorch中保存 state_dict .pth 之外也导出一个 ONNX 或 TorchScript 版本。这么做可以在出现兼容性或转换问题时有备选方案。此外，定期测试加载流程——不要假定保存成功就万事大吉，时常模拟从零加载模型跑推理或继续训练，确保文件完整性和可用性。最后，善用文档和社区资源：官方文档提供了详细的 API 用法和注意事项（例如 PyTorch 官方教程关于保存/加载模型​、TensorFlow 关于 SavedModel 的指南等），遇到问题可以查阅这些资料以获取权威解答。

结语

模型的保存和加载是深度学习工程中至关重要的一环。掌握不同格式的特点和用法，能让你在模型训练、调试和部署中游刃有余。对于开发者而言，应根据场景选择合适的格式：在开发阶段确保易用和灵活，在部署阶段确保稳定和高效。希望本文对 PyTorch 与 TensorFlow 常见模型保存格式的全面解析和代码示范，能帮助你深化对模型序列化的理解，在实际项目中避免陷坑，成功地管理好你的模型生命周期。

参考资料来源

https://pytorch.org/tutorials/beginner/saving_loading_models.html
https://jamesmccaffrey.wordpress.com/2020/09/14/saving-and-using-a-pytorch-model-in-onnx-format/
https://www.tensorflow.org/guide/saved_model
https://www.omi.me/blogs/tensorflow-guides/what-is-the-difference-between-hdf5-and-savedmodel-in-tensorflow
https://www.metriccoders.com/post/understanding-tensorflow-lite-tflite-format
https://stackoverflow.com/questions