TensorFlow微调DistilBERT保存后加载调用predict报错

问题描述 #

我是TensorFlow和BERT初学者，参考网络教程使用自有数据集微调DistilBERT完成自定义文本分类任务，数据集仅包含message和label两列，结构可参考示例图pic1。
模型训练过程顺利完成，训练阶段自定义的predict_proba函数可正常运行输出结果，但在保存模型时出现警告，保存完成后分别尝试使用keras.models.load_model()、tf.saved_model.load()、tf.keras.models.load_model三种方式加载模型，再次调用predict函数时均触发报错。

保存阶段警告 #

WARNING:tensorflow:Skipping full serialization of Keras layer <keras.layers.core.dropout.Dropout object at 0x000001B910694D88>, because it is not built.
WARNING:tensorflow:Skipping full serialization of Keras layer <keras.layers.core.dropout.Dropout object at 0x000001B97BBC58C8>, because it is not built.
WARNING:tensorflow:Skipping full serialization of Keras layer <keras.layers.core.dropout.Dropout object at 0x000001B97BCF0E48>, because it is not built.
WARNING:tensorflow:Skipping full serialization of Keras layer <keras.layers.core.dropout.Dropout object at 0x000001B91071AB08>, because it is not built.
WARNING:tensorflow:Skipping full serialization of Keras layer <keras.layers.core.dropout.Dropout object at 0x000001B91072E388>, because it is not built.
WARNING:tensorflow:Skipping full serialization of Keras layer <keras.layers.core.dropout.Dropout object at 0x000001B91073FC48>, because it is not built.
WARNING:absl:Found untraced functions such as embeddings_layer_call_fn, embeddings_layer_call_and_return_conditional_losses, transformer_layer_call_fn, transformer_layer_call_and_return_conditional_losses, LayerNorm_layer_call_fn while saving (showing 5 of 164). These functions will not be directly callable after loading.

加载后推理报错 #

核心错误为调用tf_distil_bert_for_sequence_classification层时无法找到从SavedModel加载的匹配具体函数，模型保存时记录的输入规格为shape=(None,5)的input_ids，实际传入了shape=(None,100)的input_ids与attention_mask参数，规格不匹配，完整报错信息如下：

ValueError: Exception encountered when calling layer "tf_distil_bert_for_sequence_classification" (type TFDistilBertForSequenceClassification).
Could not find matching concrete function to call loaded from the SavedModel.Got:
  Positional arguments (9 total):
    * {'input_ids': <tf.Tensor 'input_ids_1:0' shape=(None, 100) dtype=int32>, 'attention_mask': <tf.Tensor 'input_ids:0' shape=(None, 100) dtype=int32>}
    * None
    * None
    * None
    * None
    * None
    * None
    * None
    * False
  Keyword arguments: {}

 Expected these arguments to match one of the following 2 option(s):

Option 1:
  Positional arguments (9 total):
    * {'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='input_ids/input_ids')}
    * None
    * None
    * None
    * None
    * None
    * None
    * None
    * False
  Keyword arguments: {}

Option 2:
  Positional arguments (9 total):
    * {'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='input_ids/input_ids')}
    * None
    * None
    * None
    * None
    * None
    * None
    * None
    * True
  Keyword arguments: {}

Call arguments received:
  • args=({'input_ids': 'tf.Tensor(shape=(None, 100), dtype=int32)', 'attention_mask': 'tf.Tensor(shape=(None, 100), dtype=int32)'},)
  • kwargs={'training': 'False'}

复现代码 #

import pandas as pd
import tensorflow as tf
import tensorflow_hub as hub
import transformers
from transformers import DistilBertTokenizer
from transformers import TFDistilBertForSequenceClassification
from transformers import TFTrainer, TFTrainingArguments

pd.set_option('display.max_colwidth', None)
BATCH_SIZE = 16
N_EPOCHS = 3

df = pd.read_csv('twitter.csv', names=["message", "label"], encoding='cp949')

X = list(df['message'])
y = list(df['label'])
y = list(pd.get_dummies(y,drop_first=True)[True])

from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size = 0.20, random_state = 0)

tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')

train_encodings = tokenizer(X_train, truncation=True, padding=True)
test_encodings = tokenizer(X_test, truncation=True, padding=True)

train_dataset = tf.data.Dataset.from_tensor_slices((
    dict(train_encodings),
    y_train
))

test_dataset = tf.data.Dataset.from_tensor_slices((
    dict(test_encodings),
    y_test
))

model = TFDistilBertForSequenceClassification.from_pretrained("distilbert-base-uncased")

#chose the optimizer
optimizerr = tf.keras.optimizers.Adam(learning_rate=5e-5)

#define the loss function 
losss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

#build the model
model.compile(optimizer=optimizerr,
              loss=losss,
              metrics=['accuracy'])

history = model.fit(train_dataset.shuffle(len(X_train)).batch(BATCH_SIZE),
          epochs=N_EPOCHS,
          batch_size=BATCH_SIZE)

# model evaluation on the test set
model.evaluate(test_dataset.shuffle(len(X_test)).batch(BATCH_SIZE), 
               return_dict=True, 
               batch_size=BATCH_SIZE)

# tests
def predict_proba(text_list, model, tokenizer):  
    #tokenize the text
    encodings = tokenizer(text_list, 
                          max_length=1000, 
                          truncation=True, 
                          padding=True)
    #transform to tf.Dataset
    dataset = tf.data.Dataset.from_tensor_slices((dict(encodings)))
    #predict
    preds = model.predict(dataset.batch(1)).logits  
    
    #transform to array with probabilities
    res = tf.nn.softmax(preds, axis=1).numpy()      
    
    return res

examples = [
    'In 2008, several failing banks were bailed out partially using taxpayer money. Putting all money at bank provide risk, risk of devaluation, risk of inflation, risk of aggressive centralise policy. Decentralized system like bitcoin working on blockchain provide relief.',
    'Bitcoin is counterfeit. Disagree? Look again.',
    'Did I make a bad GPU purchase before the end of Ethereum mining?',
    '@Mamooetz Help. I created this bot to reply to ETH, BITCOIN, and NFT but I dont know how to shut it off.',
    '@WaldorickWilson Cryptocurrency doesnt have to be cryptic. Luno takes the complexity out of #Bitcoin and lets you buy, store, learn and earn all in one place',
]

result = predict_proba(examples, model, tokenizer)
print(result)

# save model
dataset_name = 'adv'
saved_model_path = './{}_bert'.format(dataset_name.replace('/', '_'))

# model.save(saved_model_path)
tf.saved_model.save(model, saved_model_path)

# load model
loaded_model = tf.saved_model.load(saved_model_path)
inference_function = loaded_model.signatures['serving_default']

reloaded = tf.keras.models.load_model(saved_model_path)
predict_proba(examples, reloaded, tokenizer)

问题原因 #

• 输入维度不匹配：训练阶段做tokenize时未指定固定max_length，tokenizer自动按训练集最长样本长度做padding，刚好训练集最长样本长度为5，模型保存时将输入签名固化为(None,5)的固定长度，后续推理传入长度100的输入自然无法匹配。
• 保存方式错误：直接使用tf.saved_model.save保存HuggingFace的TF模型，会导致部分未在训练阶段被激活的层（比如代码里的多个Dropout层）没有完成构建，大量内部函数未被追踪，保存后的模型结构不完整。
• tokenize参数不统一：训练阶段padding参数用了默认值按批次最长样本填充，推理阶段predict_proba里设置了max_length=1000，前后输入规则不一致。

修复方案 #

• 统一固定输入长度：所有tokenize步骤统一指定固定的max_length（根据文本长度可选128、256等常用值），padding统一设置为max_length，不要留空让tokenizer自动适配数据集长度。
• 优先使用HuggingFace原生保存/加载接口：不要直接用tensorflow自带的save方法保存Transformer类模型，用模型自带的save_pretrained方法保存，同时保存tokenizer，加载时也用对应的from_pretrained方法，能完全避免层未构建、函数丢失的问题。
• 若必须导出为TF原生SavedModel格式用于部署，保存前先传入一组符合固定输入维度的dummy数据跑一次前向传播，确保所有层完成构建后再保存。

修正后的关键代码 #

# 定义全局固定最大长度，所有tokenize步骤统一使用
MAX_LEN = 128

# 训练集、测试集编码时指定固定长度和padding规则
train_encodings = tokenizer(X_train, truncation=True, padding='max_length', max_length=MAX_LEN)
test_encodings = tokenizer(X_test, truncation=True, padding='max_length', max_length=MAX_LEN)

# 修正预测函数的tokenize参数，和训练阶段保持一致
def predict_proba(text_list, model, tokenizer):  
    encodings = tokenizer(text_list, 
                          max_length=MAX_LEN, 
                          truncation=True, 
                          padding='max_length')
    dataset = tf.data.Dataset.from_tensor_slices((dict(encodings)))
    preds = model.predict(dataset.batch(1)).logits  
    res = tf.nn.softmax(preds, axis=1).numpy()      
    return res

# 训练完成后用原生方法保存模型和分词器
model.save_pretrained(saved_model_path)
tokenizer.save_pretrained(saved_model_path)

# 加载时用对应方法
loaded_model = TFDistilBertForSequenceClassification.from_pretrained(saved_model_path)
loaded_tokenizer = DistilBertTokenizer.from_pretrained(saved_model_path)

如果必须导出为TF SavedModel格式，保存前先执行以下代码完成所有层构建：

dummy_input = {
    "input_ids": tf.constant([[0]*MAX_LEN], dtype=tf.int32),
    "attention_mask": tf.constant([[0]*MAX_LEN], dtype=tf.int32)
}
_ = model(dummy_input, training=False)
# 再执行保存操作
model.save(saved_model_path, save_format='tf')

内容的提问来源于stack exchange，提问作者pcw