TensorFlow 2.0 保存模型結構和參數

筆記摘自《Google老師親授 TensorFlow2.0 入門到進階_課程》

1. 模型保存和部署流程

TensorFlow 2.0的模型保存和部署流程如下：

2. 模型的保存

（1）保存爲HDF5

直接修改callback的參數列表如下：

# tf.keras.models.Sequential()
model = keras.models.Sequential([
    keras.layers.Flatten(input_shape=[28, 28]),
    keras.layers.Dense(300, activation='relu'),
    keras.layers.Dense(100, activation='relu'),
    keras.layers.Dense(10, activation='softmax')
])
model.compile(loss="sparse_categorical_crossentropy",
              optimizer = "sgd",
              metrics = ["accuracy"])
              
logdir = './graph_def_and_weights'
if not os.path.exists(logdir):
    os.mkdir(logdir)
output_model_file = os.path.join(logdir,
                                 "fashion_mnist_model.h5")
# 修改callback的參數列表,選擇需要保存的參數
callbacks = [
    keras.callbacks.TensorBoard(logdir),
    keras.callbacks.ModelCheckpoint(output_model_file,
                                    save_best_only = True,
                                    save_weights_only = False),
    keras.callbacks.EarlyStopping(patience=5, min_delta=1e-3),
]
# 添加參數callbacks
history = model.fit(x_train_scaled, y_train, epochs=10,
                    validation_data=(x_valid_scaled, y_valid),
                    callbacks = callbacks)   
model.evaluate(x_test_scaled, y_test, verbose=0)
# 選擇只保存weight和bias
model.save_weights(os.path.join(logdir, "fashion_mnist_weights.h5"))                

將保存的權重賦值給模型時，需要先建立一個與訓練一模一樣的模型：

# tf.keras.models.Sequential()
model = keras.models.Sequential([
    keras.layers.Flatten(input_shape=[28, 28]),
    keras.layers.Dense(300, activation='relu'),
    keras.layers.Dense(100, activation='relu'),
    keras.layers.Dense(10, activation='softmax')
])
model.compile(loss="sparse_categorical_crossentropy",
              optimizer = "sgd",
              metrics = ["accuracy"])
              
logdir = './graph_def_and_weights'
if not os.path.exists(logdir):
    os.mkdir(logdir)
output_model_file = os.path.join(logdir,
                                 "fashion_mnist_model.h5")
model.load_weights(output_model_file)
model.evaluate(x_test_scaled, y_test, verbose=0)              

（2）保存爲SaveModel

使用tf.saved_model.save()直接保存模型所有內容
保存路徑："./keras_saved_graph"

# tf.keras.models.Sequential()
model = keras.models.Sequential([
    keras.layers.Flatten(input_shape=[28, 28]),
    keras.layers.Dense(300, activation='relu'),
    keras.layers.Dense(100, activation='relu'),
    keras.layers.Dense(10, activation='softmax')
])
model.compile(loss="sparse_categorical_crossentropy",
              optimizer = "sgd",
              metrics = ["accuracy"])
history = model.fit(x_train_scaled, y_train, epochs=10,
                    validation_data=(x_valid_scaled, y_valid))
# 保存模型                  
tf.saved_model.save(model, "./keras_saved_graph")

查看保存內容：
!saved_model_cli show --dir ./keras_saved_graph --all
可以看到關於輸入和輸出的簽名信息。如下：

使用tf.saved_model.load()加載模型，驗證模型保存是否正確：

loaded_saved_model = tf.saved_model.load('./keras_saved_graph')
print(list(loaded_saved_model.signatures.keys()))
# ['serving_default']
inference = loaded_saved_model.signatures['serving_default']
print(inference.structured_outputs)
# {'dense_2': TensorSpec(shape=(None, 10), dtype=tf.float32, name='dense_2')}
results = inference(tf.constant(x_test_scaled[0:1]))
print(results['dense_2'])
# tf.Tensor([[3.3287105e-07 6.5204258e-05 7.5596186e-06 9.7430329e-06 6.4412257e-06 8.6376350e-03 2.2177779e-05 5.3723875e-02 4.3917933e-04 9.3708795e-01]], shape=(1, 10), dtype=float32)

3. 轉爲TFLite

簽名函數轉爲SavedModel、h5 model轉爲concrete_function這裏不列出，只討論如何從h5和SavedModel直接轉爲tflite model。

（1）h5轉爲TFLite model

使用tf.lite.TFLiteConverter直接將h5轉爲TFLite model。

# 加載模型
loaded_keras_model = keras.models.load_model(
    './graph_def_and_weights/fashion_mnist_model.h5')
# 轉換
keras_to_tflite_converter = tf.lite.TFLiteConverter.from_keras_model(
    loaded_keras_model)
keras_tflite = keras_to_tflite_converter.convert()
# 保存
if not os.path.exists('./tflite_models'):
    os.mkdir('./tflite_models')
with open('./tflite_models/keras_tflite', 'wb') as f:
    f.write(keras_tflite)

（2）SaveModel轉爲 tf.lite model

直接傳入 SaveModel的路徑即可。
這裏與上述第二種保存方式的保存目錄相對應。

saved_model_to_tflite_converter = tf.lite.TFLiteConverter.from_saved_model('./keras_saved_graph/')
saved_model_tflite = saved_model_to_tflite_converter.convert()
with open('./tflite_models/saved_model_tflite', 'wb') as f:
    f.write(saved_model_tflite)

4. 量化TFLite模型

量化的意思是把32位精度變爲8位，簡化模型。
直接給convetor設置一個optimization即可。

（1）量化h5模型

loaded_keras_model = keras.models.load_model(
    './graph_def_and_weights/fashion_mnist_model.h5')
keras_to_tflite_converter = tf.lite.TFLiteConverter.from_keras_model(
    loaded_keras_model)
# 量化
keras_to_tflite_converter.optimizations = [
    tf.lite.Optimize.OPTIMIZE_FOR_SIZE]
keras_tflite = keras_to_tflite_converter.convert()
if not os.path.exists('./tflite_models'):
    os.mkdir('./tflite_models')
with open('./tflite_models/quantized_keras_tflite', 'wb') as f:
    f.write(keras_tflite)

（2）量化SaveModel

saved_model_to_tflite_converter = tf.lite.TFLiteConverter.from_saved_model('./keras_saved_graph/')
saved_model_to_tflite_converter.optimizations = [
    tf.lite.Optimize.OPTIMIZE_FOR_SIZE]
saved_model_tflite = saved_model_to_tflite_converter.convert()
with open('./tflite_models/quantized_saved_model_tflite', 'wb') as f:
    f.write(saved_model_tflite)

4. 模型的部署

TensorFlow Lite 示例應用