ValueError：未知层：TFBertModel。请确保将此对象传递给“custom_objects”参数

Question

我在这里训练bert模型。下面的代码是我用来训练的，当我加载保存的模型进行预测时，它显示了这个错误。有人可以帮我吗？

import tensorflow as tf
import logging
from tensorflow.keras.layers import (
Dense,
Flatten,
Conv1D,
Dropout,
Input,)
from tensorflow.keras.models import Sequential
from tensorflow.keras.optimizers import Adam
from tensorflow.keras import Model
from tensorflow.keras import regularizers
from transformers import BertTokenizer, TFBertModel
import pandas as pd
from sklearn.model_selection import train_test_split
import numpy as np
from tqdm import tqdm
tqdm.pandas()
import re
import random 


try:
   tpu = tf.distribute.cluster_resolver.TPUClusterResolver()
   tf.config.experimental_connect_to_cluster(tpu)
   tf.tpu.experimental.initialize_tpu_system(tpu)
   strategy = tf.distribute.experimental.TPUStrategy(tpu)
except ValueError:
   strategy = tf.distribute.get_strategy() # for CPU and single GPU
   print('Number of replicas:', strategy.num_replicas_in_sync)


max_length = 512
batch_size = 16
dev_size = 0.1
num_class = 4

model_name = "bert-base-multilingual-cased"
tokenizer = BertTokenizer.from_pretrained(model_name)

train_df = pd.read_csv(path)
train, dev = train_test_split(train_df, test_size=dev_size, random_state=42)

def bert_encode(data):
    tokens = tokenizer.batch_encode_plus(
      data, max_length=max_length, padding="max_length", truncation=True
     )
    return tf.constant(tokens["input_ids"])

train_encoded = bert_encode(train.TweetText)
dev_encoded = bert_encode(dev.TweetText)

train_labels = tf.keras.utils.to_categorical(train.label.values, num_classes=num_class)
dev_labels = tf.keras.utils.to_categorical(dev.label.values, num_classes=num_class)

train_dataset = (
tf.data.Dataset.from_tensor_slices((train_encoded, train_labels))
.shuffle(100)
.batch(batch_size)
).cache()

dev_dataset = (
   tf.data.Dataset.from_tensor_slices((dev_encoded, dev_labels))
   .shuffle(100)
   .batch(batch_size)
).cache()

def bert_tweets_model():
    bert_encoder = TFBertModel.from_pretrained(model_name, output_attentions=True)
    input_word_ids = Input(
    shape=(max_length,), dtype=tf.int32, name="input_ids"
    )
    last_hidden_states = bert_encoder(input_word_ids)[0]
    clf_output = Flatten()(last_hidden_states)
    net = Dense(512, activation="relu")(clf_output)
    net = Dropout(0.3)(net)
    net = Dense(440, activation="relu")(net)
    net = Dropout(0.3)(net)
    output = Dense(num_class, activation="softplus")(net)
    model = Model(inputs=input_word_ids, outputs=output)
    return model

with strategy.scope():
   model = bert_tweets_model()
   adam_optimizer = Adam(learning_rate=1e-5)
   model.compile(
   loss="categorical_crossentropy", optimizer=adam_optimizer, metrics=["accuracy"]
   )
   model.summary()

history = model.fit(
 train_dataset,
 batch_size=batch_size,
 epochs=3,
 validation_data=dev_dataset,
 verbose=1,
 )

model.save('bert_model.h5', overwrite=True)

loaded_model = tf.keras.models.load_model('bert_model.h5', 
                                  custom_objects={"CustomModel": bert_tweets_model})

Answer 1

它告诉你它不知道什么TFBertModel是层，只需按照建议进行即可：

import transformers
loaded_model = tf.keras.models.load_model('bert_model.h5', custom_objects={"TFBertModel": transformers.TFBertModel})