tensorflow - 保存和加载使用拥抱面更改器(mutator)模型作为其第一层的 tensorflow 模型的问题

Question

您好，我在保存和加载 tensorflow 模型时遇到了一些严重的问题，该模型结合了拥抱面更改器(mutator)和一些自定义层来进行分类。我正在使用最新的 Huggingface 变形金刚 tensorflow keras 版本。思路是使用distilbert提取特征，然后通过CNN运行特征进行分类和提取。就获得正确的分类而言，我已经做好了一切工作。

问题在于训练后保存模型然后再次加载模型。

我正在使用 tensorflow keras 和 tensorflow 2.2 版

下面是设计模型、训练模型、评估模型、保存加载的代码

    bert_config = DistilBertConfig(dropout=0.2, attention_dropout=0.2, output_hidden_states=False)
    bert_config.output_hidden_states = False
    transformer_model = TFDistilBertModel.from_pretrained(DISTIL_BERT, config=bert_config)

    input_ids_in = tf.keras.layers.Input(shape=(BERT_LENGTH,), name='input_token', dtype='int32')
    input_masks_in = tf.keras.layers.Input(shape=(BERT_LENGTH,), name='masked_token', dtype='int32')

    embedding_layer = transformer_model(input_ids_in, attention_mask=input_masks_in)[0]
    x = tf.keras.layers.Bidirectional(
        tf.keras.layers.LSTM(50, return_sequences=True, dropout=0.1,
                             recurrent_dropout=0, recurrent_activation="sigmoid",
                             unroll=False, use_bias=True, activation="tanh"))(embedding_layer)
    x = tf.keras.layers.GlobalMaxPool1D()(x)

    outputs = []
    # lots of code here to define the dense layers to generate the outputs
    # .....
    # .....

    model = Model(inputs=[input_ids_in, input_masks_in], outputs=outputs)
    for model_layer in model.layers[:3]:
        logger.info(f"Setting layer {model_layer.name} to not trainable")
        model_layer.trainable = False
    rms_optimizer = RMSprop(learning_rate=0.001)
    model.compile(loss=SigmoidFocalCrossEntropy(), optimizer=rms_optimizer)

    # the code to fit the model (which works)
    # then code to evaluate the model (which also works)

    # finally saving the model. This too works.
    tf.keras.models.save_model(model, save_url, overwrite=True, include_optimizer=True, save_format="tf")

但是，当我尝试使用以下方法加载保存的模型时

    tf.keras.models.load_model(
            path, custom_objects={"Addons>SigmoidFocalCrossEntropy": SigmoidFocalCrossEntropy})

我收到以下加载错误

ValueError: The two structures don't have the same nested structure.

First structure: type=TensorSpec str=TensorSpec(shape=(None, 128), dtype=tf.int32, name='inputs')

Second structure: type=dict str={'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='inputs/input_ids')}

More specifically: Substructure "type=dict str={'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='inputs/input_ids')}" is a sequence, while substructure "type=TensorSpec str=TensorSpec(shape=(None, 128), dtype=tf.int32, name='inputs')" is not
Entire first structure:
.
Entire second structure:
{'input_ids': .}

我认为问题是因为可以使用来自 DistilBertTokenizer.encode() 的字典输入来调用 TFDistilBertModel 层，而这恰好是第一层。因此，加载时的模型编译器希望它成为调用模型的输入签名。但是，定义给模型的输入是两个形状为 (None, 128) 的张量

那么我该如何告诉加载函数或保存函数采用正确的签名呢？

Answer 1

我解决了这个问题。

问题是上面代码中的对象 transformer_model 本身不是一个层。因此，如果我们想将它嵌入到另一个 keras 层中，我们应该使用包裹在模型中的内部 keras 层

所以换行

embedding_layer = transformer_model(input_ids_in, attention_mask=input_masks_in[0]

到

embedding_layer = transformer_model.distilbert([input_ids_in, input_masks_in])[0]

让一切正常。希望这可以帮助别人。花了很长时间通过 tf.keras 代码调试来弄清楚这一点，尽管事后看来这是显而易见的。 :)