TensorFlow Serving部署Keras BERT模型REST API预测报错求助

预测错误输出 #

{
"error": "Op type not registered 'TFText>RoundRobinTrim' in binary running on ljh-my-keras-bert-model. Make sure the Op and Kernel are registered in the binary running in this process. Note that if you are loading a saved graph which used ops from tf.contrib (e.g. tf.contrib.resampler), accessing should be done before importing the graph, as contrib ops are lazily registered when the module is first accessed."
}

本地环境版本 #

Python 3.10.14
tensorflow                                   2.18.0
tensorflow-datasets                          4.9.6
tensorflow-io-gcs-filesystem                 0.37.1
tensorflow-metadata                          1.16.1
tensorflow-text                              2.18.0
keras                                        3.6.0
keras-hub-nightly                            0.16.1.dev202410210343
keras-nlp                                    0.17.0

模型定义代码 #

import os
os.environ["KERAS_BACKEND"] = "tensorflow"  # "jax" or "tensorflow" or "torch"

import tensorflow_datasets as tfds
import keras_nlp

imdb_train, imdb_test = tfds.load(
    "imdb_reviews",
    split=["train", "test"],
    as_supervised=True,
    batch_size=16,
)

import keras
# 加载模型
classifier = keras_nlp.models.BertClassifier.from_preset(
    "bert_tiny_en_uncased",
    num_classes=2,
    activation="softmax",
)
# 编译模型
classifier.compile(
    loss="sparse_categorical_crossentropy",
    optimizer=keras.optimizers.Adam(5e-5),
    metrics=["sparse_categorical_accuracy"],
    jit_compile=True,
)
# 微调
classifier.fit(imdb_train.take(250), validation_data=imdb_test.take(250))
# 预测示例
classifier.predict(["What an amazing movie!", "A total waste of my time."])
# 预期输出: array([[0.34156954, 0.65843046], [0.52648497, 0.473515  ]], dtype=float32)

模型保存代码 #

import tensorflow as tf
import keras_nlp

def preprocess(inputs):
    # 将输入字符串转换为token ID、padding mask和segment ID
    preprocessor = classifier.preprocessor
    encoded = preprocessor(inputs)
    return {
        'token_ids': encoded['token_ids'],
        'padding_mask': encoded['padding_mask'],
        'segment_ids': encoded['segment_ids']
    }

@tf.function(input_signature=[tf.TensorSpec(shape=[None], dtype=tf.string)])
def serving_fn(inputs):
    preprocessed = preprocess(inputs)
    outputs = classifier(preprocessed)
    return outputs

# 保存模型
model_export_path = "/Users/jiahao.liu/tf_saved_models/my-keras-bert-model/1"
tf.saved_model.save(
    classifier,
    export_dir=model_export_path,
    signatures={"serving_default": serving_fn}
)

print(f"模型已保存至: {model_export_path}")

Docker镜像构建文件 #

FROM tensorflow/serving:latest

COPY my-keras-bert-model /models/my_keras_bert_model

# 设置模型名称环境变量
ENV MODEL_NAME my_keras_bert_model
# ENV OMP_NUM_THREADS 8
# ENV TF_NUM_INTEROP_THREADS 8
# ENV TF_NUM_INTRAOP_THREADS 8

# 启动TensorFlow Serving
CMD tensorflow_model_server --port=8500 --rest_api_port=8501 --model_name=${MODEL_NAME} --model_base_path=/models/${MODEL_NAME}

预测请求 #

POST http://localhost:8501/v1/models/my_keras_bert_model/versions/1:predict
Content-Type: application/json

{"instances": ["What an amazing movie!", "A total waste of my time."]}

解决方法 #

报错核心原因是TensorFlow Serving官方默认镜像未包含tensorflow-text库，而你的模型推理流程中用到了TFText的RoundRobinTrim操作，导致运行时找不到对应的Op定义。

方案1：自定义包含tensorflow-text的Serving镜像 #

修改Dockerfile，在官方镜像基础上安装与本地环境版本一致的tensorflow-text：

FROM tensorflow/serving:latest

# 安装对应版本的tensorflow-text
RUN pip install tensorflow-text==2.18.0

COPY my-keras-bert-model /models/my_keras_bert_model

ENV MODEL_NAME my_keras_bert_model

CMD tensorflow_model_server --port=8500 --rest_api_port=8501 --model_name=${MODEL_NAME} --model_base_path=/models/${MODEL_NAME}

方案2：导出端到端整合模型 #

将预处理逻辑与模型推理打包成一个完整的SavedModel，确保所有依赖操作都被固化：

# 替换原有的保存逻辑，创建端到端模型
class EndToEndModel(tf.keras.Model):
    def __init__(self, classifier):
        super().__init__()
        self.classifier = classifier
        self.preprocessor = classifier.preprocessor
    
    def call(self, inputs):
        preprocessed = self.preprocessor(inputs)
        return self.classifier(preprocessed)

# 实例化端到端模型
end_to_end_model = EndToEndModel(classifier)

# 保存完整模型
tf.saved_model.save(
    end_to_end_model,
    export_dir=model_export_path,
    signatures={
        "serving_default": tf.function(
            lambda inputs: end_to_end_model(inputs),
            input_signature=[tf.TensorSpec(shape=[None], dtype=tf.string)]
        )
    }
)

之后使用方案1的自定义镜像部署即可。

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