使用拥抱面变压器模型作为第一层的保存和加载张量流模型的问题

Question

您好，我在保存和加载张量流模型时遇到一些严重的问题，该模型是拥抱面部变压器+一些自定义层进行分类的组合。我正在使用最新的 Huggingface Transformers 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': .}

我认为问题是因为 TFDistilBertModel 层可以使用来自 DistilBertTokenizer.encode() 的字典输入来调用，而这恰好是第一层。因此，加载时的模型编译器期望这是调用模型的输入签名。然而，定义到模型的输入是两个形状张量（无，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 代码进行调试才弄清楚这一点，尽管事后看来这是显而易见的。:)