TensorFlow 2.0 模型保存与转换：.pb格式详解

版本要求：

Tensorflow 2.0.0+

tensorflow原生保存的模型不适用于多平台使用，将模型转化为.pb格式，可以更加方便的转化为别的格式，本文主要介绍如何转化.pb格式。

以mnist数据集为例进行讲解：

1. 创建网络模型，训练，保存模型：

inputs = tf.keras.Input(shape=(28,28,1), name='input')
# [28, 28, 1] => [28, 28, 64]
input = tf.keras.layers.Flatten(name="flatten")(inputs)
fc_1 = tf.keras.layers.Dense(512, activation='relu', name='fc_1')(input)
fc_2 = tf.keras.layers.Dense(256, activation='relu', name='fc_2')(fc_1)
pred = tf.keras.layers.Dense(10, activation='softmax', name='output')(fc_2)

model = tf.keras.Model(inputs=inputs, outputs=pred, name='mnist')
model.summary()

将网络搭建出来后，进行训练，然后保存模型，训练过程不再赘述，训练完保存模型，有两种保存方式：

方式1：

from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
import tensorflow as tf

# training code -----------------
tf.saved_model.save(obj=model, export_dir="./model")
# Convert Keras model to ConcreteFunction
# 注意这个Input，是自己定义的输入层名
full_model = tf.function(lambda Input: model(Input))
full_model = full_model.get_concrete_function(
    tf.TensorSpec(model.inputs[0].shape, model.inputs[0].dtype))

# Get frozen ConcreteFunction
frozen_func = convert_variables_to_constants_v2(full_model)
frozen_func.graph.as_graph_def()

layers = [op.name for op in frozen_func.graph.get_operations()]
print("-" * 50)
print("Frozen model layers: ")
for layer in layers:
    print(layer)

print("-" * 50)
print("Frozen model inputs: ")
print