如何正确加载并使用已保存的BERT机器学习模型？

你遇到的错误本质是输入格式不匹配：训练时用text.Transformer把文本转成了模型需要的张量（input_ids、attention_mask这类整数序列），但加载模型后直接传原始字符串，模型没法处理这种输入。下面针对三种保存格式，分别给出正确的使用方法：

方法一：加载.tf完整模型 #

加载模型后，必须用和训练时一样的text.Transformer实例预处理文本，再把处理后的张量传入模型：

import tensorflow as tf
import pandas as pd
# 替换成你实际导入text模块的路径
from your_script import text  

df = pd.read_csv('dataset.csv', delimiter=';')
df = df.drop(['first label', 'second label', 'third label'], axis=1)
X = df['post'].values

# 初始化和训练时完全相同的Transformer预处理工具
MODEL_NAME = 'bert-large-uncased'
t = text.Transformer(MODEL_NAME)
# 对要预测的文本做预处理，得到模型需要的张量格式
processed_input = t.preprocess_test(X[:1], verbose=True)

# 加载模型
model = tf.keras.models.load_model('mymodel.tf')
# 直接传入预处理后的字典格式数据
prediction = model.predict(processed_input)[0][0]
print(prediction)

方法二：加载.h5权重文件 #

这种方式需要先重构模型结构，再加载权重，同样要做预处理：

import tensorflow as tf
import pandas as pd
from your_script import text  

df = pd.read_csv('dataset.csv', delimiter=';')
df = df.drop(['first label', 'second label', 'third label'], axis=1)
X = df['post'].values

# 重构和训练时一样的模型结构
MODEL_NAME = 'bert-large-uncased'
t = text.Transformer(MODEL_NAME)
model = t.get_classifier()
# 加载预训练好的权重
model.load_weights("myweights.h5")

# 预处理文本
processed_input = t.preprocess_test(X[:1], verbose=True)
prediction = model.predict(processed_input)[0][0]
print(prediction)

方法三：加载SavedModel格式（my_saved_model文件夹） #

和.tf模型类似，需要预处理文本，然后调用模型的签名函数：

import tensorflow as tf
import pandas as pd
from your_script import text  

df = pd.read_csv('dataset.csv', delimiter=';')
df = df.drop(['first label', 'second label', 'third label'], axis=1)
X = df['post'].values

MODEL_NAME = 'bert-large-uncased'
t = text.Transformer(MODEL_NAME)
processed_input = t.preprocess_test(X[:1], verbose=True)

# 加载SavedModel
model = tf.saved_model.load('my_saved_model')
# 获取默认的推理签名函数
infer_func = model.signatures['serving_default']
# 传入预处理后的张量，键名要匹配模型输入
output = infer_func(**processed_input)
# 根据模型输出的键名获取结果，一般是'logits'或'predictions'
prediction = output['logits'].numpy()[0][0]
print(prediction)

关键提醒 #

• 预处理必须和训练时完全一致：用相同的模型名称、分词器配置，否则肯定会出现格式不匹配。
• BERT这类Transformer模型天生就不能直接处理原始字符串，必须转成整数张量才能输入。
• 如果你是自己写的text.Transformer模块，确保导入路径正确，别出现找不到模块的问题。

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