浅谈tensorflow模型保存为pb的各种姿势

浅谈TensorFlow模型保存为pb的各种姿势

在TensorFlow中，我们可以将训练好的模型保存为pb文件，以便在其他地方使用。本文将浅谈TensorFlow模型保存为pb的各种姿势，并提供两个示例说明。

方法1：使用tf.saved_model.save()保存模型

在TensorFlow 2.0中，我们可以使用tf.saved_model.save()方法将模型保存为pb文件。以下是保存模型的示例代码：

import tensorflow as tf

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

# 编译模型
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

# 训练模型
model.fit(x_train, y_train, epochs=10, validation_data=(x_test, y_test))

# 保存模型
tf.saved_model.save(model, 'saved_model')

在这个示例中，我们使用tf.saved_model.save()方法将模型保存为pb文件，并将其保存在'saved_model'文件夹中。

方法2：使用tf.compat.v1.saved_model.simple_save()保存模型

在TensorFlow 1.x中，我们可以使用tf.compat.v1.saved_model.simple_save()方法将模型保存为pb文件。以下是保存模型的示例代码：

import tensorflow as tf

# 构建模型
x = tf.placeholder(tf.float32, [None, 784])
W = tf.Variable(tf.zeros([784, 10]))
b = tf.Variable(tf.zeros([10]))
y = tf.nn.softmax(tf.matmul(x, W) + b)

# 保存模型
with tf.compat.v1.Session() as sess:
    sess.run(tf.compat.v1.global_variables_initializer())
    tf.compat.v1.saved_model.simple_save(sess, 'saved_model', inputs={'x': x}, outputs={'y': y})

在这个示例中，我们使用tf.compat.v1.saved_model.simple_save()方法将模型保存为pb文件，并将其保存在'saved_model'文件夹中。

示例1：使用tf.saved_model.save()保存模型

以下是使用tf.saved_model.save()方法保存模型的示例代码：

import tensorflow as tf

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

# 编译模型
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

# 训练模型
model.fit(x_train, y_train, epochs=10, validation_data=(x_test, y_test))

# 保存模型
tf.saved_model.save(model, 'saved_model')

在这个示例中，我们使用tf.saved_model.save()方法将模型保存为pb文件，并将其保存在'saved_model'文件夹中。

示例2：使用tf.compat.v1.saved_model.simple_save()保存模型

以下是使用tf.compat.v1.saved_model.simple_save()方法保存模型的示例代码：

import tensorflow as tf

# 构建模型
x = tf.placeholder(tf.float32, [None, 784])
W = tf.Variable(tf.zeros([784, 10]))
b = tf.Variable(tf.zeros([10]))
y = tf.nn.softmax(tf.matmul(x, W) + b)

# 保存模型
with tf.compat.v1.Session() as sess:
    sess.run(tf.compat.v1.global_variables_initializer())
    tf.compat.v1.saved_model.simple_save(sess, 'saved_model', inputs={'x': x}, outputs={'y': y})

在这个示例中，我们使用tf.compat.v1.saved_model.simple_save()方法将模型保存为pb文件，并将其保存在'saved_model'文件夹中。

结语

以上是TensorFlow模型保存为pb的各种姿势的完整攻略，包括使用tf.saved_model.save()和tf.compat.v1.saved_model.simple_save()两种方法，并提供了两个示例说明。在实际应用中，我们可以根据具体情况来选择合适的方法来保存模型。