tensorflow - 无法加载 SavedModel

Question

我正在使用 tf.keras 训练一些模型，并希望保存训练后的模型。有两种推荐的方法，tfSavedModel和 keras.h5文件。然而，事情变得非常混乱SavedModel。

以下是一些重现该问题的简短脚本：

import numpy as np
import tensorflow as tf

def my_network():
    backbone_model = tf.keras.applications.InceptionResNetV2(
        include_top=False, input_shape=(224, 224, 3), weights="imagenet", pooling="avg"
    )

    inputs = tf.keras.layers.Input(shape=(224, 224, 3), name="images")
    backbone_features = backbone_model(inputs)
    pre_embeddings = tf.keras.layers.Dense(
        512,
        activation=None,
        name="pre_embeddings",
        kernel_regularizer=tf.keras.regularizers.l2(),
    )(backbone_features)

    embeddings = tf.keras.layers.Lambda(
        lambda x: tf.math.l2_normalize(x, 1, 1e-10), name="embeddings"
    )(pre_embeddings)

    probs = tf.keras.layers.Dense(
        1000,
        activation="softmax",
        name="predictions",
        kernel_regularizer=tf.keras.regularizers.l2(),
        bias_regularizer=tf.keras.regularizers.l2(),
    )(pre_embeddings)

    return tf.keras.Model(inputs, [embeddings, probs], name="my_network")


img_arr = np.random.rand(1, 224, 224, 3)

resnet_model = my_network()
emb_1, _ = resnet_model.predict(img_arr)
resnet_model.save("./resnet_model.h5")

new_model = tf.keras.models.load_model('./resnet_model.h5')
emb_2, _ = new_model.predict(img_arr)

np.testing.assert_array_almost_equal(emb_1, emb_2)

上面的脚本可以正常工作。但是，当我尝试以SavedModel格式保存时（通过.h5从模型路径中删除），它失败了。模型保存成功，但加载时报错，报错信息为：

NotImplementedError: When subclassing the `Model` class, you should implement a `call` method.

我很困惑，因为我没有使用任何子类模型。如脚本所示，我的网络仅使用功能 API 构建。

标签： tensorflowkerastensorflow2.0