使用Kashgari模型部署TensorFlow Serving服务指南-CSDN博客

使用Kashgari模型部署TensorFlow Serving服务指南

前言

TensorFlow Serving是一个高性能的机器学习模型服务系统，专为生产环境设计。本文将详细介绍如何将Kashgari训练好的NLP模型部署为TensorFlow Serving服务，并提供完整的客户端调用示例。

准备工作

在开始之前，请确保已安装以下环境：

TensorFlow Serving (推荐使用Docker方式)
Kashgari最新版本
Python请求库requests

模型保存与转换

首先需要将训练好的Kashgari模型转换为TensorFlow Serving可识别的格式：

from kashgari.tasks.classification import BiGRU_Model
from kashgari.corpus import SMP2018ECDTCorpus
from kashgari import utils

# 加载训练数据
train_x, train_y = SMP2018ECDTCorpus.load_data()

# 训练模型
model = BiGRU_Model()
model.fit(train_x, train_y)

# 转换为SavedModel格式
utils.convert_to_saved_model(model,
                             model_path="saved_model/bgru",
                             version=1)

转换后的模型会保存在指定目录下，包含模型结构和权重信息。

启动TensorFlow Serving服务

使用Docker启动TensorFlow Serving服务是最便捷的方式：

docker run -t --rm -p 8501:8501 \
  -v "<绝对路径>/saved_model:/models/" \
  -e MODEL_NAME=bgru \
  tensorflow/serving

参数说明：

-p 8501:8501：将容器的8501端口映射到主机
-v：挂载模型目录
-e MODEL_NAME：指定模型名称

客户端调用模型服务

1. 获取模型元数据

在调用前，我们需要了解模型的输入输出结构：

import requests
res = requests.get("http://localhost:8501/v1/models/bgru/metadata")
inputs = res.json()['metadata']['signature_def']['signature_def']['serving_default']['inputs']
input_sample_keys = list(inputs.keys())
print(input_sample_keys)

根据返回结果，我们会看到不同的输入键名，这决定了后续调用时的数据格式。

2. 单输入字段调用（非BERT模型）

对于传统的BiGRU等模型，通常只有一个输入字段：

import requests
import numpy as np
from kashgari.processors import load_processors_from_model

# 加载预处理工具
text_processor, label_processor = load_processors_from_model('模型路径')

# 准备样本数据
samples = [
    ['hello', 'world'],
    ['你', '好', '世', '界']
]

# 文本预处理
tensor = text_processor.transform(samples)
instances = [i.tolist() for i in tensor]

# 调用预测接口
r = requests.post("http://localhost:8501/v1/models/bgru:predict", 
                 json={"instances": instances})
predictions = r.json()['predictions']

# 转换预测结果为标签
labels = label_processor.inverse_transform(np.array(predictions).argmax(-1))
print(labels)

3. 多输入字段调用（BERT模型）

对于BERT等需要多输入的模型，调用方式有所不同：

import requests
import numpy as np
from kashgari.processors import load_processors_from_model

# 加载预处理工具
text_processor, label_processor = load_processors_from_model('模型路径')

# 准备样本数据
samples = [
    ['hello', 'world'],
    ['你', '好', '世', '界']
]

# 文本预处理
tensor = text_processor.transform(samples)

# 构建多字段输入
instances = [{
   "Input-Token": i.tolist(),
   "Input-Segment": np.zeros(i.shape).tolist()
} for i in tensor]

# 调用预测接口
r = requests.post("http://localhost:8501/v1/models/bgru:predict", 
                 json={"instances": instances})
predictions = r.json()['predictions']

# 转换预测结果为标签
labels = label_processor.inverse_transform(np.array(predictions).argmax(-1))
print(labels)