tinybert基于openvino服务化部署
- 部署请求流程：客户端（直接发送文本text） -> 服务端（tokenizer，条件分支, 模型推理，softmax分类）-> 客户端（接收返回结果）
- 区别：该部署方法与tinybert-mediapipe的部署方法基本一样，唯一的区别就是在预处理（preprocess）阶段添加了一个条件分支，当输入满足条件时，跳过模型推理，直接输入到后处理（postprocess）阶段。
- 完整的代码项目请参考：tinybert-openvino
部署前准备
文件夹结构
与tinybert-mediapipe的文件夹结构完整一样，只有部分文件需要更改：
- config.json
- preprocess.py
- postprocess.py
- graph.pbtxt
配置文件
- config.json
json
{
"model_config_list": [
{
"config": {
"name": "tinybert",
"base_path": "tinybert-mediapipe-route/models",
"target_device": "CPU",
"model_version_policy": {
"latest": {"num_versions": 1}
},
"nireq": 4,
"plugin_config": {
"PERFORMANCE_HINT": "THROUGHPUT"
},
"shape": {
"input_ids": "(?,100)",
"attention_mask": "(?,100)"
}
}
}
],
"mediapipe_config_list": [
{
"name": "tinybert-mediapipe-route",
"graph_path": "tinybert-mediapipe-route/graph.pbtxt"
}
]
}其实配置文件config.json基本不变，只是改了name和base_path而已，其它保持不变
预处理脚本
- preprocess.py
此处预处理脚本需要增加条件分支内容，当条件满足时，跳过模型推理，直接输入到后处理阶段。
python
from transformers import AutoTokenizer
import numpy as np
from pyovms import Tensor
import traceback
import sys
class OvmsPythonModel:
def initialize(self, kwargs):
# 分词器路径硬编码为相对于图根目录的路径
self.tokenizer = AutoTokenizer.from_pretrained("tinybert-mediapipe-route/tokenizer")
self.max_length = 100
self.inference = True
def truncation_head_tail(self, text):
# 截取输入文本的前50个字符和后50个字符
if len(text) <= self.max_length:
return text
half = self.max_length // 2
return text[:half] + "..." + text[-half:]
def execute(self, inputs):
try:
# 输入是 pyovms.Tensor 列表，取第一个
text_tensor = inputs[0]
# 将数据解码为字符串
# text = bytes(text_tensor).decode("utf-8")
if hasattr(text_tensor, 'to_string'):
text = text_tensor.to_string()
elif hasattr(text_tensor, 'get_bytes'):
text = text_tensor.get_bytes().decode("utf-8")
else:
# 兜底方案：如果必须用 bytes()，请确保只取有效载荷部分
# 注意：这取决于具体版本的序列化格式，仅作调试参考
raw = bytes(text_tensor)
# 尝试寻找第一个有效的 UTF-8 字符起始位置，或根据协议跳过固定头部
text = raw.decode("utf-8", errors="ignore").strip('\x00')
# 分词
encoded = self.tokenizer(
self.truncation_head_tail(text),
return_tensors="np",
truncation=True,
padding="max_length",
max_length=self.max_length,
)
input_ids = encoded["input_ids"].astype(np.int64)
attention_mask = encoded["attention_mask"].astype(np.int64)
token_type_ids = encoded["token_type_ids"].astype(np.int64)
## 条件分支内容，例如通过bm2.5判断句子相似性，服务则设置self.inference=False，跳过模型推理，直接把结果输送给后处理
self.inference = False
predicted_class = 1
confidence = 0.95
if self.inference:
# 创建输出 Tensor，名称必须与输出流一致
out_ids = Tensor("input_ids_py", input_ids, datatype="INT64")
out_mask = Tensor("attention_mask_py", attention_mask, datatype="INT64")
out_type_ids = Tensor("token_type_ids_py", token_type_ids, datatype="INT64")
return [out_ids, out_mask, out_type_ids]
else:
predicted_class = np.array(predicted_class, dtype=np.int64)
out_predicted_class = Tensor("predicted_class_py", predicted_class, datatype="INT64")
# 如果predicted_class的值为字符串
# predicted_class = predicted_class.encode("utf-8")
# out_predicted_class = Tensor("predicted_class_py", predicted_class, datatype="STRING")
confidence = np.array(confidence, dtype=np.float32)
out_confidence = Tensor("confidence_py", confidence, datatype="FP32")
return [out_predicted_class, out_confidence]
except Exception as e:
# 打印完整堆栈到 stderr，将被 OVMS 日志捕获
traceback.print_exc(file=sys.stderr)
# raise # 继续抛出，保持图错误行为
# ✅ 返回一个全 0 的合法 Tensor 作为兜底，让后处理节点去处理
dummy_ids = np.zeros((1, self.max_length), dtype=np.int64)
dummy_mask = np.zeros((1, self.max_length), dtype=np.int64)
dummy_type_ids = np.zeros((1, self.max_length), dtype=np.int64)
return [
Tensor("input_ids_py", dummy_ids, datatype="INT64"),
Tensor("attention_mask_py", dummy_mask, datatype="INT64"),
Tensor("token_type_ids_py", dummy_type_ids, datatype="INT64")
]后处理脚本
- postprocess.py
此处需要判断接收的输入是模型输出(last_hidden_state_py)还是条件分支输出(predicted_class_py, confidence_py)
python
import json
import numpy as np
from pyovms import Tensor
class OvmsPythonModel:
@staticmethod
def softmax(x):
e_x = np.exp(x - np.max(x, axis=-1, keepdims=True))
return e_x / e_x.sum(axis=-1, keepdims=True)
def execute(self, inputs):
# 1. 检查条件判断的输出（直接来自预处理）
predicted_calss_tensor = None
confidence_tensor = None
for t in inputs:
if t.name == "predicted_calss_py":
predicted_calss_tensor = t
elif t.name == "confidence_py":
confidence_tensor = t
if predicted_calss_tensor is not None and confidence_tensor is not None:
# 句子相似度分支：直接使用预处理的结果
# predicted_calss 是 INT64，转为 Python int
predicted_calss = int(np.array(predicted_calss_tensor))
# 如果 predicted_calss_tensor 是 bytes，解码为字符串
# predicted_calss = bytes(predicted_calss_tensor).decode('utf-8')
# confidence 是 FP32，转为 Python float
confidence = float(np.array(confidence_tensor))
result = {
"predicted_calss": predicted_calss,
"confidence": round(confidence, 4)
}
return [Tensor("response", json.dumps(result).encode('utf-8'))]
# 2. 推理路径：处理模型推理输出的 last_hidden_state_py
logits = None
for t in inputs:
if t.name == "last_hidden_state_py":
logits = np.array(t) # [batch, 100, 312]
if logits is None :
raise ValueError("Missing logits")
# 获取token的CLS索引
cls_logits = logits[:, 0, :] # [batch 312]
# 计算softmax得到概率分布
cls_probs = self.softmax(cls_logits) # [batch 312]
# 获取每个样本的预测类别和置信度
predicted_class = np.argmax(cls_probs, axis=-1).astype(np.int64) # [batch,]
confidence = np.round(np.max(cls_probs, axis=-1).astype(np.float32), 4) # [batch,]
result = {
"predicted_class": f"{predicted_class}",
"confidence": round(confidence)
}
# 返回字符串
return [Tensor("response", json.dumps(result).encode('utf-8'))]计算图
- graph.pbtxt
计算图中，预处理节点需要同时把输出给模型推理节点和条件分支节点的输出都要写上，后处理节点也需要同时把模型输出的节点和条件分支节点的输出都要写上；不管走哪条分支，输入输出节点都要写上，openvino的机制是：如果节点为空，该分支就不会流通，只走有数据流通的分支。
pbtxt
input_stream: "OVMS_PY_TENSOR:text"
# 预处理
node {
name: "preprocess_node"
calculator: "PythonExecutorCalculator"
input_side_packet: "PYTHON_NODE_RESOURCES:py"
input_stream: "INPUT:text"
output_stream: "INPUT_IDS:input_ids_py"
output_stream: "ATTENTION_MASK:attention_mask_py"
output_stream: "TOKEN_TYPE_IDS:token_type_ids_py"
output_stream: "PREDICTED_CLASS:predicted_class_py" # 把条件分支作为输出
output_stream: "CONFIDENCE:confidence_py" # 把条件分支作为输出
node_options {
[type.googleapis.com/mediapipe.PythonExecutorCalculatorOptions] {
handler_path: "tinybert-mediapipe-route/preprocess.py"
}
}
}
# 转换器：PY -> OV (input_ids)
node {
calculator: "PyTensorOvTensorConverterCalculator"
input_stream: "OVMS_PY_TENSOR:input_ids_py"
output_stream: "OVTENSOR:input_ids"
}
# 转换器：PY -> OV (attention_mask)
node {
calculator: "PyTensorOvTensorConverterCalculator"
input_stream: "OVMS_PY_TENSOR:attention_mask_py"
output_stream: "OVTENSOR:attention_mask"
}
# 转换器：PY -> OV (token_type_ids)
node {
calculator: "PyTensorOvTensorConverterCalculator"
input_stream: "OVMS_PY_TENSOR:token_type_ids_py"
output_stream: "OVTENSOR:token_type_ids"
}
# ===== 推理会话（从 model_config_list 加载） =====
node {
calculator: "OpenVINOModelServerSessionCalculator"
output_side_packet: "SESSION:session"
node_options: {
[type.googleapis.com/mediapipe.OpenVINOModelServerSessionCalculatorOptions]: {
servable_name: "tinybert-doublehead"
servable_version: "1"
}
}
}
# ===== 推理节点 =====
node {
calculator: "OpenVINOInferenceCalculator"
input_side_packet: "SESSION:session"
input_stream: "OVTENSOR:input_ids"
input_stream: "OVTENSOR1:attention_mask" # 第2个输入用 OVTENSOR1
input_stream: "OVTENSOR2:token_type_ids" # 第3个输入用 OVTENSOR2
output_stream: "OVTENSOR:last_hidden_state"
node_options: {
[type.googleapis.com/mediapipe.OpenVINOInferenceCalculatorOptions]: {
tag_to_input_tensor_names {
key: "OVTENSOR"
value: "input_ids"
}
tag_to_input_tensor_names {
key: "OVTENSOR1"
value: "attention_mask"
}
tag_to_input_tensor_names {
key: "OVTENSOR2"
value: "token_type_ids"
}
tag_to_output_tensor_names {
key: "OVTENSOR"
value: "last_hidden_state"
}
}
}
}
# 转换器：OV -> PY (last_hidden_state)
node {
calculator: "PyTensorOvTensorConverterCalculator"
input_stream: "OVTENSOR:last_hidden_state"
output_stream: "OVMS_PY_TENSOR:last_hidden_state_py"
node_options {
[type.googleapis.com/mediapipe.PyTensorOvTensorConverterCalculatorOptions] {
tag_to_output_tensor_names {
key: "OVMS_PY_TENSOR"
value: "last_hidden_state_py"
}
}
}
}
# 后处理
node {
name: "postprocess_node"
calculator: "PythonExecutorCalculator"
input_side_packet: "PYTHON_NODE_RESOURCES:py"
input_stream: "INPUT:last_hidden_state_py"
input_stream: "PREDICTED_CLASS:predicted_class_py" # 把预处理的条件分支作为输入
input_stream: "CONFIDENCE:confidence_py" # 把预处理的条件分支作为输入
output_stream: "OUTPUT:response"
node_options {
[type.googleapis.com/mediapipe.PythonExecutorCalculatorOptions] {
handler_path: "tinybert-mediapipe-route/postprocess.py"
}
}
}
output_stream: "OVMS_PY_TENSOR:response"镜像
构建镜像部分和tinybert-mediapipe的一样，根据条件分支判断需要哪些依赖进行自定义安装即可。
模型部署
bash
docker run -d --name ovms \
-p 127.0.0.1:8000:8000 \
-p 127.0.0.1:8001:8001 \
-v ./tinybert-mediapip-route:/workspace/tinybert-mediapipe-route \
openvino/model_server:py3-transformers \
ovms --config_path /workspace/tinybert-mediapipe-route/config.json --port 8000 --rest_port 8001注意： 镜像名称替换为自行构建的镜像名称
服务请求
bash
python http_tinybert-mediapipe-route.py --http http://localhost:8001 --model tinybert-medaipipe-route --input "tinybert 是一个分号的文本分类器，可以识别任何文本的意图！"注意： --model的名称要和config.json中的mediapipe_config_list中的name一致, 而不是model_config_list下的name; 和tinybert-mediapipe的区别在于，tinybert-mediapipe-route的模型名称是tinybert-medaipipe-route，而tinybert-mediapipe的模型名称是tinybert-mediapipe
其中，http_tinybert-mediapipe.py文件内容如下：
python
import requests
import time
import json
import argparse
class HTTP_CLINET:
def __init__(self, http_url="http://localhost:8001", model='tinybert'):
self.url = f"{http_url}/v2/models/{model}/infer"
def post(self, text):
try:
payload = {
"inputs": [
{
"name": "text",
"shape": [1],
"datatype": "BYTES",
"data": [text]
}
]
}
headers = {
"Content-Type": "application/json",
# "Authorization": f"Bearer {self.api_key}" # 注意 Bearer 后面有一个空格
}
response = requests.post(
self.url,
json=payload,
headers=headers,
timeout=30
)
if response.status_code == 200:
result = response.json()
print("推理成功!")
# print(f"响应数据: {json.dumps(result, indent=2)}")
return result
else:
print(f"请求失败，状态码: {response.status_code}")
print(f"错误信息: {response.text}")
return None
except requests.exceptions.RequestException as e:
print(f"请求异常: {e}")
return None
def parse_opt():
parser = argparse.ArgumentParser(description="http client")
parser.add_argument('--http', '-u', type=str, default='http://localhost:8001', help="http_url")
parser.add_argument('--model', '-m', type=str, default="tinybert-mediapipe-route", help='openvino server name')
parser.add_argument('--input', '-i', type=str, default=None, help='input text')
return parser.parse_args()
def main(args):
http_client = HTTP_CLINET(http_url=args.http, model=args.model)
defeat_text = r'tinybert 是一个分号的文本分类器，可以识别任何文本的意图！'
text = args.input if args.input else defeat_text
print(f"text: {text}")
print("执行推理...")
start_time = time.time()
result = http_client.post(text=text)
end_time = time.time()
print("请求推理完成")
print(f"请求推理总耗时：{(end_time-start_time): .3f} s")
print(result)
if __name__ == "__main__":
args = parse_opt()
main(args)性能测试
bash
python performance_tinybert-mediapipe-route.py --http http://localhost:8001 --model tinybert-mediapipe- 和tinybert-mediapipe的区别在于，tinybert-mediapipe-route的模型名称是tinybert-medaipipe-route，而tinybert-mediapipe的模型名称是tinybert-mediapipe
其中，performance_tinybert-mediapipe-route.py文件内容如下：
python
import requests
import time
import os
import json
import concurrent.futures
import statistics
import random
import csv
from tqdm import tqdm
import argparse
import numpy as np
class PerformanceBenchmark:
"""高性能并发推理测试器"""
def __init__(
self,
http_url: str,
model_name: str,
text_path: str,
min_text_len: int = 100,
max_text_len: int = 100,
request_timeout: int = 30,
max_seq_len: int = 100
):
"""
:param http_url: Triton 服务基础 URL (如 http://localhost:8001)
:param model_name: 模型名称
:param tokenizer_path: 分词器路径
:param text_path: 测试文本 JSON 文件路径
:param min_text_len: 保留文本的最小长度
:param max_text_len: 文本截断到的最大长度
:param request_timeout: 单次请求超时(秒)
"""
self.http_url = http_url.rstrip("/")
self.model_name = model_name
self.url = f"{self.http_url}/v2/models/{self.model_name}/infer"
self.request_timeout = request_timeout
self.max_seq_len = max_seq_len
# 加载并预处理文本数据
self.texts = self._load_texts(text_path, min_text_len, max_text_len)
if not self.texts:
raise RuntimeError("没有可用的测试文本，请检查数据文件或长度过滤条件。")
def _load_texts(self, text_path: str, min_len: int, max_len: int) -> list:
"""加载 Alpaca 格式 JSON 并拼接为文本列表"""
combined_texts = []
try:
with open(text_path, "r", encoding="utf-8") as f:
data = json.load(f)
for item in data:
instruction = item.get("instruction", "")
input_text = item.get("input", "")
output_text = item.get("output", "")
combined = instruction + input_text + output_text
combined = combined.replace(" ", "").replace("\n", "")
if len(combined) < min_len:
continue
if len(combined) >= max_len:
combined = combined[:max_len]
combined_texts.append(combined)
return combined_texts
except FileNotFoundError:
print(f"错误: 文件 {text_path} 不存在")
except json.JSONDecodeError:
print(f"错误: 文件 {text_path} 不是有效的JSON格式")
return []
def _single_benchmark(
self,
num_requests: int,
concurrency: int,
verbose: bool = False,
) -> dict:
"""单次并发测试"""
latencies = []
errors = 0
def worker(_):
nonlocal errors
session = requests.Session()
text = random.choice(self.texts)
payload = {
"inputs": [
{
"name": "text",
"shape": [1],
"datatype": "BYTES",
"data": [text]
}
]
}
start = time.time()
try:
resp = session.post(self.url, json=payload, timeout=self.request_timeout)
elapsed = (time.time() - start) * 1000
if resp.status_code == 200:
return elapsed
else:
errors += 1
if verbose:
print(f"请求失败状态码: {resp.status_code}, 响应: {resp.text}")
print(f"text len: {len(text)}")
print(f"text: {text}")
return None
except Exception as e:
errors += 1
if verbose:
print(f"请求异常: {e}")
return None
finally:
session.close()
start_total = time.time()
with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as executor:
futures = [executor.submit(worker, i) for i in range(num_requests)]
for future in tqdm(
concurrent.futures.as_completed(futures),
total=num_requests,
desc=f"并发{concurrency}",
):
result = future.result()
if result is not None:
latencies.append(result)
end_total = time.time()
total_time = end_total - start_total
success_count = len(latencies)
if success_count == 0:
print("所有请求均失败！")
return None
qps = success_count / total_time
avg_latency = statistics.mean(latencies)
max_lat = max(latencies)
min_lat = min(latencies)
latencies_sorted = sorted(latencies)
p50 = latencies_sorted[int(len(latencies_sorted) * 0.5)]
p90 = latencies_sorted[int(len(latencies_sorted) * 0.9)]
p99 = latencies_sorted[int(len(latencies_sorted) * 0.99)]
results = {
"concurrency": concurrency,
"num_requests": num_requests,
"success_count": success_count,
"error_count": errors,
"total_time_sec": round(total_time, 2),
"qps": round(qps, 2),
"avg_latency_ms": round(avg_latency, 2),
"min_latency_ms": round(min_lat, 2),
"max_latency_ms": round(max_lat, 2),
"p50_latency_ms": round(p50, 2),
"p90_latency_ms": round(p90, 2),
"p99_latency_ms": round(p99, 2),
}
print(f"\n========== 并发数 {concurrency} 测试结果 ==========")
print(f"总请求数: {num_requests}")
print(f"成功请求数: {success_count}")
print(f"失败请求数: {errors}")
print(f"总耗时(s): {total_time:.2f}")
print(f"吞吐量(QPS): {qps:.2f}")
print(f"平均延迟(ms): {avg_latency:.2f}")
print(f"最小延迟(ms): {min_lat:.2f}")
print(f"最大延迟(ms): {max_lat:.2f}")
print(f"P50 延迟(ms): {p50:.2f}")
print(f"P90 延迟(ms): {p90:.2f}")
print(f"P99 延迟(ms): {p99:.2f}")
print("==============================")
return results
def run_benchmark_sweep(
self,
concurrency_list: list = [1, 2, 4, 8, 16, 32],
num_requests_per_test: int = 200,
output_csv: str = "benchmark_results.csv",
cooldown_sec: int = 2,
):
"""
遍历多个并发数进行测试，并保存结果至 CSV。
:param concurrency_list: 并发数列表
:param num_requests_per_test: 每次测试的总请求数
:param output_csv: 结果输出 CSV 文件路径
:param cooldown_sec: 每轮测试后的冷却时间（秒）
"""
all_results = []
for concurrency in concurrency_list:
print(f"\n========== 测试并发数: {concurrency} ==========")
result = self._single_benchmark(
num_requests=num_requests_per_test,
concurrency=concurrency,
verbose=True,
)
if result is not None:
all_results.append(result)
time.sleep(cooldown_sec)
if all_results:
keys = all_results[0].keys()
os.makedirs(os.path.dirname(output_csv) or ".", exist_ok=True)
with open(output_csv, "w", newline="") as f:
writer = csv.DictWriter(f, fieldnames=keys)
writer.writeheader()
writer.writerows(all_results)
print(f"性能指标已保存至 {output_csv}")
else:
print("没有成功的结果可保存。")
def parse_opt():
parser = argparse.ArgumentParser(description="TinyBert-MediaPipe模型基于OpenVINO的性能测试")
parser.add_argument("--http", "-u", type=str, default="http://localhost:8001", help="OpenVINO 服务基础 URL")
parser.add_argument("--model", "-m", type=str, default="tinybert-doublehead-mediapipe-route", help="模型名称")
parser.add_argument("--input", "-i", type=str, default="../datasets/alpaca_gpt4_data_zh.json", help="测试文本 JSON 文件")
parser.add_argument("--max_batch_size", type=int, default=8, help="最大批次大小（仅用于文件名标识）")
parser.add_argument("--concurrency", "-c", type=str, default="1,4,8,16", help="并发数列表，逗号分隔")
parser.add_argument("--num", "-n", type=int, default=128, help="每次测试的总请求数")
parser.add_argument("--output_dir", "-o", type=str, default="performance_result", help="输出目录")
parser.add_argument("--max_seq_len", type=int, default=100, help="最大输入文本长度")
return parser.parse_args()
def main():
args = parse_opt()
output_dir = args.output_dir
os.makedirs(output_dir, exist_ok=True)
output_csv = os.path.join(
output_dir,
f"{args.model}_cpu_openvino_batch{args.max_batch_size}_seqlen{args.max_seq_len}_benchmark_results.csv",
)
concurrency_list = [int(x.strip()) for x in args.concurrency.split(",")]
# 初始化测试器
benchmark = PerformanceBenchmark(
http_url=args.http,
model_name=args.model,
text_path=args.input,
min_text_len=128, # 可根据需要调整
max_text_len=1024, # 固定文本截断长度，与模型序列长度一致
max_seq_len=args.max_seq_len
)
# 执行测试
benchmark.run_benchmark_sweep(
concurrency_list=concurrency_list,
num_requests_per_test=args.num,
output_csv=output_csv,
cooldown_sec=2,
)
if __name__ == "__main__":
main()