TensorFlow从入门到实践:TensorFlow模型保存与加载入门

TensorFlow模型保存与加载实战

学习目标

通过本课程的学习，学员将掌握在TensorFlow中如何保存和加载模型的方法，包括模型的权重、结构以及整个模型的保存与加载。还将了解如何在不同的环境中复用模型，以提高开发效率和模型的可移植性。

相关知识点

• TensorFlow模型保存加载与复用

学习内容

1 TensorFlow模型保存加载与复用

1.1 TensorFlow模型保存

TensorFlow 提供了高效的模型保存功能，帮助开发者将训练好的模型持久化到磁盘，以便后续加载与复用。通过 model.save() 方法，用户可以将模型的结构、权重、优化器状态及训练配置等信息统一保存为 .h5 或 SavedModel 格式文件。例如，model.save(‘model_weights.weights.h5’) 会生成一个包含所有必要信息的文件，便于后续直接使用。加载模型时，通过 tf.keras.models.load_model() 函数即可恢复模型状态，如 loaded_model = tf.keras.models.load_model(‘my_model.h5’)，加载后的模型与原始模型在结构和权重上完全一致，可直接用于推理或继续训练。其中，SavedModel 格式支持跨平台兼容性，并包含完整的计算图和元数据，适合生产环境部署。保存后的文件可脱离原始代码独立存在，方便团队协作或模型共享。这一机制不仅简化了模型版本管理，还为迁移学习、模型部署等场景提供了便捷支持，是 TensorFlow 模型生命周期管理的核心功能之一。

在深度学习项目中，模型训练是一个耗时且资源密集的过程。因此，能够有效地保存训练好的模型对于后续的模型测试、部署和复用至关重要。TensorFlow提供了多种方法来保存模型，包括保存模型的权重、结构以及整个模型。

保存模型的权重

保存模型的权重是最基本的保存方式，它只保存了模型训练过程中学习到的参数值。这种方式适用于需要重新定义模型结构的情况，即在加载模型时需要重新构建模型的结构。

%pip install tensorflow

import tensorflow as tf
import numpy as np
from tensorflow.keras import layers, models

# 定义一个简单的卷积神经网络
def create_model():
    model = models.Sequential([
        layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),
        layers.MaxPooling2D((2, 2)),
        layers.Conv2D(64, (3, 3), activation='relu'),
        layers.MaxPooling2D((2, 2)),
        layers.Flatten(),
        layers.Dense(64, activation='relu'),
        layers.Dense(10, activation='softmax')
    ])
    model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])
    return model

# 创建模型
model = create_model()

# 假设模型已经训练完成

# 保存模型的权重
model.save_weights('model_weights.weights.h5')  # 修改文件名以符合要求

保存模型的结构

保存模型的结构意味着保存了模型的架构信息，但不包括模型的权重。这种方式适用于需要保存模型结构但不保存权重的情况，例如在不同的环境中重新训练模型。

# 保存模型的结构
json_config = model.to_json()
with open('model_config.json', 'w') as json_file:
    json_file.write(json_config)

保存整个模型

保存整个模型是最全面的保存方式，它不仅保存了模型的权重和结构，还包括了优化器的状态、损失函数等信息。这种方式适用于需要完全复现训练过程的情况。

# 保存整个模型
model.save('complete_model.h5')

1.2 TensorFlow模型加载

TensorFlow 的模型加载功能让开发者能够快速恢复已保存的模型状态，实现无缝复用。通过 tf.keras.models.load_model() 函数，用户可直接加载 .h5 或 SavedModel 格式的模型文件，系统会自动解析其结构、权重和训练配置。加载后的模型保留原始功能，可直接用于推理预测或继续训练。SavedModel 格式因其包含完整的计算图和元数据，支持在异构环境（如从开发环境迁移到生产服务器）中保持一致性。

加载模型是模型保存的逆过程，它允许在不同的环境中复用已经训练好的模型。根据保存的方式不同，加载模型的方法也有所不同。

加载模型的权重

加载模型的权重需要先定义模型的结构，然后将保存的权重加载到模型中。

# 重新创建模型
model = create_model()

# 加载模型的权重
model.load_weights('model_weights.weights.h5')

加载模型的结构

加载模型的结构需要先读取保存的结构信息，然后根据结构信息重新构建模型。

# 读取模型的结构
with open('model_config.json', 'r') as json_file:
    json_config = json_file.read()

# 重新构建模型
model = tf.keras.models.model_from_json(json_config)

加载整个模型

加载整个模型是最简单的方式，它可以直接从文件中加载模型的所有信息。

# 加载整个模型
model = tf.keras.models.load_model('complete_model.h5')

1.3 模型的复用

TensorFlow 的模型复用机制通过保存与加载的标准化流程，显著提升了开发效率与应用灵活性。加载后的模型不仅完整保留原始结构和权重，还可直接应用于新数据集的推理任务，或作为预训练模型参与迁移学习。例如，开发者可在加载的模型基础上添加新层、调整超参数，快速适配不同业务场景。这种复用方式避免了重复训练的资源消耗，同时支持跨项目共享模型资产。结合 TensorFlow 的跨平台兼容性，复用后的模型可无缝部署到边缘设备、云端服务器或移动端，为快速迭代产品功能、缩短交付周期提供了有力支持。

在不同的硬件平台上运行模型

模型可以在不同的硬件平台上运行，例如从CPU到GPU，或者从本地机器到云服务器。这需要确保模型的保存和加载方式与目标平台兼容。

# 假设有一个测试数据集，形状为 (num_samples, 28, 28, 1)
# 这里使用随机数据作为示例
num_samples = 1000  # 测试样本数量
x_test = np.random.rand(num_samples, 28, 28, 1).astype('float32')  # 生成随机测试数据
 
# 加载模型
model = tf.keras.models.load_model('complete_model.h5')
 
# 在 GPU 上进行预测(或者直接在CPU上运行，直接写predictions = model.predict(x_test))
with tf.device('/GPU:0'):
    predictions = model.predict(x_test)
# `predictions` 是模型的输出，通常是形状为 (num_samples, num_classes) 的数组
print(predictions.shape)

在不同的项目中使用模型

在不同的项目中使用模型可以避免重复训练，提高开发效率。这需要确保模型的输入输出格式与新项目的需求一致。

# 在新项目中加载模型
model = tf.keras.models.load_model('complete_model.h5')
num_samples = 10  # 新数据样本数量
new_data = np.random.rand(num_samples, 28, 28, 1).astype('float32')  # 生成随机新数据

# 使用模型进行预测
predictions = model.predict(new_data)

在不同的数据集上测试模型

在不同的数据集上测试模型可以评估模型的泛化能力。这需要确保新数据集的格式与模型的输入格式一致。

# 在新数据集上测试模型
model = tf.keras.models.load_model('complete_model.h5')

# 这里使用随机整数作为示例
new_labels = np.random.randint(0, 10, size=(num_samples,))  # 生成随机标签，假设有10个类


# 评估模型
loss, accuracy = model.evaluate(new_data, new_labels)
print(f'Loss: {loss}, Accuracy: {accuracy}')