TensorFlow2.0模型保存与加载：ckpt、h5、pb详解

"本教程主要涵盖了TensorFlow2.0中的模型保存与加载，包括ckpt、h5和pb三种格式，并提供了相关代码示例。" 在TensorFlow2.0中，模型的保存和加载是深度学习项目中不可或缺的环节，它使得我们可以持久化训练结果，以便在后续使用或继续训练时快速恢复模型状态。本教程主要涉及了三种模型保存格式：ckpt、h5和pb，以及相应的保存和加载方法。 1. ckpt格式 ckpt（checkpoint）格式是TensorFlow中用于保存模型参数的标准方式。它由三个文件组成： - `checkpoint` 文件：记录当前模型的最新状态，包含指向数据文件和索引文件的引用。 - `data` 文件：存储模型的变量值，即张量。 - `index` 文件：提供对`data`文件中数据的索引，方便快速定位和恢复模型参数。 2. h5格式 h5（HDF5）格式是Keras框架常用的模型保存格式，它将模型的结构和参数整合到一个单一的文件中，便于整体保存和加载。相比于ckpt，h5更便于跨平台和跨库操作。 3. pb格式 pb（Protocol Buffers）格式是Google开发的一种序列化协议，适用于模型的部署和跨语言交互。在TensorFlow中，pb文件通常用于生产环境，因为它封装了模型的结构和参数，具有语言独立性，可以在不同系统和框架间无缝迁移。 模型保存 - ckpt格式保存：可以使用`tf.train.Checkpoint`对象来保存模型，例如： ```python tf.train.Checkpoint(model=model).save('model_path') ``` - h5格式保存：对于Keras模型，可以使用`save`或`save_weights`方法： ```python model.save('model.h5') # 保存整个模型（结构+权重） model.save_weights('weights.h5') # 只保存权重，不保存结构 ``` - pb格式保存：可以使用`tf.keras.models.save_model`来生成pb文件： ```python tf.keras.models.save_model(model, 'model.pb', save_format='tf') ``` 模型加载 - ckpt格式加载：利用`tf.train.Checkpoint`加载模型： ```python new_model = MyModel() tf.train.Checkpoint(model=new_model).restore(tf.train.latest_checkpoint('model_path')) ``` - h5格式加载：Keras模型提供了方便的加载函数： ```python model = tf.keras.models.load_model('model.h5') ``` - pb格式加载：加载pb模型通常涉及构建一个与原始模型结构匹配的新模型，然后用pb文件中的数据填充： ```python with tf.gfile.GFile('model.pb', 'rb') as f: graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) _ = tf.import_graph_def(graph_def, name='') input_tensor = tf.get_default_graph().get_tensor_by_name('input:0') output_tensor = tf.get_default_graph().get_tensor_by_name('output:0') ``` 在实际应用中，选择哪种格式取决于具体需求。如果目标是继续训练或评估，ckpt和h5格式可能更合适；而如果是为了部署，pb格式通常更优，因为它可以直接在生产环境中运行，无需额外的解析步骤。了解并掌握这些保存和加载方法，能够帮助我们更有效地管理和使用TensorFlow模型。