tensorflow 模型的保存与加载

tensorflow 模型的保存与加载

代码

# -*- coding: utf-8 -*-
"""
Created on  2020/11/20 9:41
@Author: CY
@email: [email protected]
"""

# pip install -q pyyaml h5py  # 以 HDF5 格式保存模型所必须
import os

import tensorflow as tf
from tensorflow import keras

print(tf.version.VERSION)

print("# 1 数据集")
(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()
train_labels = train_labels[:1000]
test_labels = test_labels[:1000]
train_images = train_images[:1000].reshape(-1, 28 * 28) / 255.0
test_images = test_images[:1000].reshape(-1, 28 * 28) / 255.0

print("# 2 模型")
# 定义一个简单的序列模型
def create_model():
    model = tf.keras.models.Sequential([
        keras.layers.Dense(512, activation='relu', input_shape=(784,)),
        keras.layers.Dropout(0.2),
        keras.layers.Dense(10)
    ])

    model.compile(optimizer='adam',
                  loss=tf.losses.SparseCategoricalCrossentropy(from_logits=True),
                  metrics=['accuracy'])

    return model


# 创建一个基本的模型实例
model = create_model()
# 显示模型的结构
model.summary()

print("#3. 在训练期间保存模型（以 checkpoints 形式保存）")
# 您可以使用训练好的模型而无需从头开始重新训练，或在您打断的地方开始训练，以防止训练过程没有保存。
# tf.keras.callbacks.ModelCheckpoint 允许在训练的过程中和结束时回调保存的模型。
# Checkpoint 回调用法
# 创建一个只在训练期间保存权重的 tf.keras.callbacks.ModelCheckpoint 回调：
checkpoint_path = "tmp/training_1/cp.ckpt"
checkpoint_dir = os.path.dirname(checkpoint_path)

# 创建一个保存模型权重的回调
cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_path,
                                                 save_weights_only=True,
                                                 verbose=1)

# 使用新的回调训练模型
model.fit(train_images,
          train_labels,
          epochs=10,
          validation_data=(test_images, test_labels),
          callbacks=[cp_callback])  # 通过回调训练

# 这可能会生成与保存优化程序状态相关的警告。
# 这些警告（以及整个笔记本中的类似警告）
# 是防止过时使用，可以忽略。


print("#4 建立一个未训练的模型")
# 创建一个新的未经训练的模型。仅恢复模型的权重时，必须具有与原始模型具有相同网络结构的模型。
# 由于模型具有相同的结构，您可以共享权重，尽管它是模型的不同实例。
# 现在重建一个新的未经训练的模型，并在测试集上进行评估。未经训练的模型将在机会水平（chance levels）上执行（准确度约为10％）：
# 创建一个基本模型实例
model = create_model()

# 评估模型
loss, acc = model.evaluate(test_images, test_labels, verbose=2)
print("Untrained model, accuracy: {:5.2f}%".format(100 * acc))
print("#4. 加载权重...")
# 加载权重
model.load_weights(checkpoint_path)

# 重新评估模型
loss, acc = model.evaluate(test_images, test_labels, verbose=2)
print("Restored model, accuracy: {:5.2f}%".format(100 * acc))

# checkpoint 回调选项
## 为 checkpoint 提供唯一名称并调整 checkpoint 频率。
## 训练一个新模型，每五个 epochs 保存一次唯一命名的 checkpoint ：

# 在文件名中包含 epoch (使用 `str.format`)
checkpoint_path = "tmp/training_2/cp-{epoch:04d}.ckpt"
checkpoint_dir = os.path.dirname(checkpoint_path)

# 创建一个回调，每 5 个 epochs 保存模型的权重
cp_callback = tf.keras.callbacks.ModelCheckpoint(
    filepath=checkpoint_path,
    verbose=1,
    save_weights_only=True,
    period=5)

# 创建一个新的模型实例
model = create_model()

# 使用 `checkpoint_path` 格式保存权重
model.save_weights(checkpoint_path.format(epoch=0))

# 使用新的回调训练模型
model.fit(train_images,
          train_labels,
          epochs=50,
          callbacks=[cp_callback],
          validation_data=(test_images, test_labels),
          verbose=0)

## 默认的 tensorflow 格式仅保存最近的5个 checkpoint 。
## 如果要进行测试，请重置模型并加载最新的 checkpoint ：
latest = tf.train.latest_checkpoint(checkpoint_dir)
# 创建一个新的模型实例
model = create_model()

# 加载以前保存的权重
model.load_weights(latest)

# 重新评估模型
loss, acc = model.evaluate(test_images, test_labels, verbose=2)
print("Restored model, accuracy: {:5.2f}%".format(100 * acc))

## 文件内容：
#### 代码将权重存储到 checkpoint—— 格式化文件的集合中，这些文件仅包含二进制格式的训练权重
#### Checkpoints 包含
######## 一个或多个包含模型权重的分片。
######## 索引文件，指示哪些权重存储在哪个分片中。
### 如果你只在一台机器上训练一个模型，你将有一个带有后缀的碎片： .data-00000-of-00001

print("# 5.手动保存权重")
# 使用 Model.save_weights 方法手动保存它们
# 默认情况下， tf.keras 和 save_weights 特别使用 TensorFlow checkpoints 格式 .ckpt 扩展名和 ( 保存在 HDF5 扩展名为 .h5 保存并序列化模型 )：
# 保存权重
model.save_weights('./tmp/checkpoints/my_checkpoint')
# 创建模型实例
model = create_model()
# 恢复权重
model.load_weights('./tmp/checkpoints/my_checkpoint')
# 评估模型
loss, acc = model.evaluate(test_images, test_labels, verbose=2)
print("Restored model, accuracy: {:5.2f}%".format(100 * acc))

print("#6 保存整个模型")
# 调用 model.save 将保存模型的结构，权重和训练配置保存在单个文件/文件夹中。
# 这可以让您导出模型，以便在不访问原始 Python 代码*的情况下使用它。
# 因为优化器状态（optimizer-state）已经恢复，您可以从中断的位置恢复训练。
# 整个模型可以以两种不同的文件格式（SavedModel 和 HDF5）进行保存。
# 需要注意的是 TensorFlow 的 SavedModel 格式是 TF2.x. 中的默认文件格式。
# 但是，模型仍可以以 HDF5 格式保存。下面介绍了以两种文件格式保存整个模型的更多详细信息。
# 保存完整模型会非常有用——您可以在 TensorFlow.js（Saved Model, HDF5）加载它们，
# 然后在 web 浏览器中训练和运行它们，或者使用 TensorFlow Lite 将它们转换为在移动设备上运行（Saved Model, HDF5）
# *自定义对象（例如，子类化模型或层）在保存和加载时需要特别注意。请参阅下面的保存自定义对象部分

# SavedModel 格式
## SavedModel 格式是序列化模型的另一种方法。以这种格式保存的模型，可以使用 tf.keras.models.load_model 还原，
## 并且模型与 TensorFlow Serving 兼容。SavedModel 指南详细介绍了如何提供/检查 SavedModel。以下部分说明了保存和还原模型的步骤。
# 创建并训练一个新的模型实例。
model = create_model()
model.fit(train_images, train_labels, epochs=5)

# 将整个模型另存为 SavedModel。
# !mkdir -p saved_model
# SavedModel 格式是一个包含 protobuf 二进制文件和 Tensorflow 检查点（checkpoint）的目录
model.save('tmp/saved_model/my_model')

print("#7. 从保存的模型重新加载新的 Keras 模型")
new_model = tf.keras.models.load_model('tmp/saved_model/my_model')
# 检查其架构
new_model.summary()
# 评估还原的模型
loss, acc = new_model.evaluate(test_images,  test_labels, verbose=2)
print('Restored model, accuracy: {:5.2f}%'.format(100*acc))

print(new_model.predict(test_images).shape)

print("#8 使用HDF5 格式 保存 模型 恢复模型")
# 创建并训练一个新的模型实例
model = create_model()
model.fit(train_images, train_labels, epochs=5)

# 将整个模型保存为 HDF5 文件。
# '.h5' 扩展名指示应将模型保存到 HDF5。
model.save('tmp/my_model.h5')
# 重新创建完全相同的模型，包括其权重和优化程序
new_model = tf.keras.models.load_model('tmp/my_model.h5')

# 显示模型的结构
new_model.summary()