博客

TensorFlow模型保存/载入的两种方法

  • Post category:Python

好的，下面是关于“TensorFlow模型保存/载入的两种方法”的完整攻略。

1. TensorFlow模型保存/载入方法

在TensorFlow中，我们可以两种方法来保存和载入模型：SavedModel和checkpoint。SavedModel是一种标准的TensorFlow型格式，可以跨平台使用，而checkpoint则是一种TensorFlow特有的模型格式，只能在TensorFlow中使用。

1.1 SavedModel方法

SavedModel是一种标准的TensorFlow模型格式，可以跨平台使用。以下是一个简单的Python程序，可以使用SavedModel方法保存和载入模型：

import tensorflow as tf

# 构建模型
model = tf.keras.Sequential([
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(10)
])

# 编译模型
model.compile(optimizer=tf.keras.optimizers.Adam(0.01),
              loss=tf.keras.losses.CategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

# 训练模型
model.fit(x_train, y_train, epochs=5)

# 保存模型
model.save('my_model')

# 载入模型
new_model = tf.keras.models.load_model('my_model')

在上面的代码中，我们首先构建了一个简单的神经网络模型。然后，我们使用compile()函数编译模型，并使用fit()函数训练模型。接着，我们使用save()函数保存模型，并使用load_model()函数载入模型。

1.2 checkpoint方法

checkpoint是一种TensorFlow特有的模型格式，只能在TensorFlow中使用。以下是一个简单的Python程序，可以使用checkpoint方法保存和载入模型：

import tensorflow as tf

# 构建模型
model = tf.keras.Sequential([
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(10)
])

# 编译模型
model.compile(optimizer=tf.keras.optimizers.Adam(0.01),
              loss=tf.keras.losses.CategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

# 训练模型
model.fit(x_train, y_train, epochs=5)

# 保存模型
checkpoint_path = "training/cp.ckpt"
cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_path,
                                                 save_weights_only=True,
                                                 verbose=1)

model.save_weights(checkpoint_path)

# 载入模型
new_model = tf.keras.Sequential([
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(10)
])

new_model.compile(optimizer=tf.keras.optimizers.Adam(0.01),
                  loss=tf.keras.losses.CategoricalCrossentropy(from_logits=True),
                  metrics=['accuracy'])

new_model.load_weights(checkpoint_path)

在上面的代码中，我们首先构建了一个简单的神经网络模型。然后，我们使用compile()函数编译模型，并使用fit()函数训练模型。接着，我们使用ModelCheckpoint()函数设置checkpoint路径，并使用save_weights()函数保存模型。最后，我们使用load_weights()函数载入模型。

2. 结语

本文介绍了TensorFlow模型保存/载入的两种方法：SavedModel和checkpoint。如果您需要跨平台使用模型，SavedModel是一个非常好的选择。如果您只需要在TensorFlow中使用模型，checkpoint是一个非常好的选择。

Tags: