带有张量流的拥抱脸转换器将两个文件保存为模型权重

Posted 2023-03-29

【中文标题】带有张量流的拥抱脸转换器将两个文件保存为模型权重

【英文标题】：huggingface transformer with tensorflow saves two files as model weights

【发布时间】：2021-09-25 23:22:25

【问题描述】：

这是我为分类任务构建模型的方式：

    def bert_for_classification(transformer_model_name, max_sequence_length, num_labels):
        config = ElectraConfig.from_pretrained(
            transformer_model_name,
            num_labels=num_labels,
            output_hidden_states=False,
            output_attentions=False
        )
        model = TFElectraForSequenceClassification.from_pretrained(transformer_model_name, config=config)
        # This is the input for the tokens themselves(words from the dataset after encoding):
        input_ids = tf.keras.layers.Input(shape=(max_sequence_length,), dtype=tf.int32, name='input_ids')

        # attention_mask - is a binary mask which tells BERT which tokens to attend and which not to attend.
        # Encoder will add the 0 tokens to the some sequence which smaller than MAX_SEQUENCE_LENGTH,
        # and attention_mask, in this case, tells BERT where is the token from the original data and where is 0 pad
        # token:
        attention_mask = tf.keras.layers.Input((max_sequence_length,), dtype=tf.int32, name='attention_mask')

        # Use previous inputs as BERT inputs:
        output = model([input_ids, attention_mask])[0]
        output = tf.keras.layers.Dense(num_labels, activation='softmax')(output)
        model = tf.keras.models.Model(inputs=[input_ids, attention_mask], outputs=output)

        model.compile(loss=keras.losses.CategoricalCrossentropy(),
                      optimizer=keras.optimizers.Adam(3e-05, epsilon=1e-08),
                      metrics=['accuracy'])

        return model

训练完这个模型后，我使用model.save_weights('model.hd5') 保存它 但事实证明有两个文件被保存：model.hd5.index 和 model.hd5.data-00000-of-00001

我应该如何从磁盘加载这个模型？

【问题讨论】：

你安装hd5py包了吗？

@Kaveh 是的，我已经安装了h5py

哦。你写了扩展hd5，而它是h5。 model.save_weights('model.h5')，如果不是h5格式，则保存为SavedModel格式。

【参考方案1】：

您有 2 种可能性来保存模型，以 keras h5 格式或 tensorflow SavedModel 格式。

您可以通过传递save_format 参数来确定格式，并将其设置为"h5" 或"tf"。如果您不指定此参数，则格式将由您传递的名称确定。如果名称有.h5后缀，则保存在keras中，否则为SavedModel格式。

不管怎样，既然你指定了后缀hd5而不是h5，它会以SavedModel格式保存。

您可以像保存它们一样简单地加载它们。

model.save_weights("model.h5")   #h5 format
model.load_weights("model.h5")
#or
model.save_weights("mymodel")    #SavedModel format
model.load_weights("mymodel")