Keras保存模型的各种方法

Keras训练模型有多种保存方法，可以保存为hdf5文件，也可以保存为json格式文件，可以同时保存模型图和权重，也可以单独保存模型图和权重，还可以保存为tensorflow-serving支持的pb格式。下面以一个简单的模型分别来介绍不同的保存方法。

模型图构建

下面用keras中函数式API构建一个简单的LSTM多分类模型，模型具体结构如下：

import keras
from keras.models import Sequential
from keras.layers import Input, Dense, Dropout, Embedding, LSTM
from keras.models import Model
from keras import backend as K
K.clear_session()

maxlen = 50
vocab_size = 1024
embedding_size = 128
hidden_size = 128
num_classes = 10

inputs = Input(shape=(maxlen,), name="inputs")
x = Embedding(input_dim=vocab_size, output_dim=embedding_size)(inputs)
x = LSTM(units=hidden_size)(x)
x = Dropout(rate=0.5)(x)
outputs = Dense(num_classes, activation='softmax', name="outputs")(x)
model = Model(inputs=inputs, outputs=outputs)

print(model.summary())

模型结构图如下：

训练模型

这里只是一个训练示例，为了减小数据预处理的过程，我们直接用随机产生的虚拟数据。

# 生成虚拟数据
import numpy as np

data = np.random.randint(vocab_size, size=(1000, maxlen))
labels = np.random.randint(num_classes, size=(1000, 1))
onehot_labels = keras.utils.to_categorical(labels, num_classes=10)

x_train, y_train = data[:800], onehot_labels[:800]
x_test, y_test = data[800:], onehot_labels[800:]

model.compile(loss='categorical_crossentropy',
              optimizer='rmsprop',
              metrics=['accuracy'])
model.fit(x_train, y_train, batch_size=32, epochs=10)
score, acc = model.evaluate(x_test, y_test, batch_size=32)
print('Test score: %s, Test acc: %s' % (score, acc))

保存模型

下面分别介绍保存为hdf5格式，json格式，以及tf-serving中的pb格式。

hdf5格式持久化

keras模型保存为hdf5又分为同时保存模型图和权重，只保存权重不保存模型图

一、同时保存模型图和权重

from keras.models import load_model

# 保存训练好的model为hdf5文件
model.save('model/my_model.h5')  
# 重新加载模型
model = load_model('model/my_model.h5')

二、只保存权重不保存模型图

# 只保存模型权重
model.save_weights('model/my_model_weights.h5')

# 加载模型权重，这里加载权重前需要重新构建model的模型图
model = Model(inputs=..., outputs=...)
model.load_weights('model/my_model_weights.h5')

json格式持久化

我们尝试把模型保存为json格式，模型图和权重可以单独保存，也可以同时保存。

一、json格式模型图

from keras.models import model_from_json

# 模型图保存为json格式字符串
model_json = model.to_json()
# 从json字符串加载模型
model = model_from_json(model_json)

如果需要持久化保存加入json的保存和读取操作

from keras.models import model_from_json

model_json = model.to_json()

# json格式模型图保存到文件
json_dict = json.loads(model_json)
with open("model/model_graph.json", "w") as fw:
    json.dump(json_dict, fw, indent=4)
    
# 加载json格式模型图
with open("model/model_graph.json", "r") as fr:
    param_data = json.load(fr)
model_json = json.dumps(param_data)

model = model_from_json(model_json)

二、json格式保存权重
keras中并没有直接把权重保存成json格式的方法，但是有获取参数名和权重的方法，调用model.get_weights()返回模型中所有权重张量的列表，类型为 Numpy 数组。

1. 简单的方法，直接调用get_weights和set_weights方法。

# 获取模型权重
weights = model.get_weights() 
# 给模型权重赋值，注意列表中的数组必须与 get_weights() 返回的权重具有相同的尺寸。
model.set_weights(weights)

2. 持久化的方法，没有详细测试，可能有些模型不行。model.layers可以获得模型的所有层，layer.weights可以获得层内参数的信息，layer.get_weights()可以获得层内每个参数的权重值。

import json

# 获得模型参数和对应权重
weights_dict = {}
for layer in model.layers:
    for weight, value in zip(layer.weights, layer.get_weights()):
        weight_value = value.tolist()
        weights_dict[weight.name] = weight_value

# 模型权重保存为json格式
with open("model/model_weights.json", "w") as fw:
    json.dump(weights_dict, fw, indent=4)

tf-serving调用的pb格式

keras模型也可以转成tf-serving调用的pb格式，利用tf-serving加载模型可以同时支持grpc和Restful API调用。模型转化方法如下：

import os
import tensorflow as tf
from keras import backend as K
from keras.models import load_model
from keras.models import Model

def export_model(model, export_model_dir, model_version):
    with tf.get_default_graph().as_default():
        # prediction_signature
        tensor_info_input = tf.saved_model.utils.build_tensor_info(model.input)
        tensor_info_output = tf.saved_model.utils.build_tensor_info(model.output)
        
        prediction_signature = (
            tf.saved_model.signature_def_utils.build_signature_def(
                inputs={'inputs': tensor_info_input}, # Tensorflow.TensorInfo
                outputs={'outpus': tensor_info_output},
                method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME)
        )

        # set-up a builder
        os.mkdir(export_model_dir)
        export_path = os.path.join(tf.compat.as_bytes(export_model_dir), tf.compat.as_bytes(str(model_version)))
        builder = tf.saved_model.builder.SavedModelBuilder(export_path)
        builder.add_meta_graph_and_variables(
            sess=K.get_session(),
            tags=[tf.saved_model.tag_constants.SERVING],
            signature_def_map={
                'predict':prediction_signature,
                tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: prediction_signature,
            },
            )
        builder.save()

model = load_model("model/model_save.h5")
export_model(model, "test_model", "1")

参看文档：
[1]. https://keras.io/zh/models/about-keras-models/
[2]. https://keras.io/zh/getting-started/faq/
[3]. https://www.tensorflow.org/tfx/tutorials/serving/rest_simple
[4]. https://www.tensorflow.org/tfx/serving/serving_basic