TensorRT 模型推論服務器

一個基於 TensorRT 的高性能深度學習模型推論服務器，支援 Embedding、Reranker 和 NLI 模型的快速推論。

功能特色

高性能推論: 基於 NVIDIA TensorRT 優化的模型推論
動態組批處理: 支援自動批次聚合，提升推論效率
多模型支援: 同時支援 Embedding、Reranker 和 NLI 模型
RESTful API: 提供標準的 HTTP API 接口
OpenAI 相容: 支援 OpenAI SDK 格式的 API 調用
GPU 記憶體優化: 高效的 GPU 記憶體管理

支援的模型

目前僅支援 Embedding、Reranker、NLI

安裝與設定

1. 克隆專案

git clone https://github.com/FubonDS/TensorrtServer.git
cd TensorrtServer

2. 安裝依賴

可參考docker/docker-compose.yaml作為基礎鏡像

pip install -r requirements.txt

轉換模型

模型轉換分為兩個步驟：PyTorch → ONNX 和 ONNX → TensorRT。

步驟一：轉換為 ONNX

相關腳本位於 trt_convert/ 目錄，均支援 argparse 參數傳入。

Embedding 模型（BGE-M3）

python trt_convert/embedding2onnx.py \
  --model_path ./embedding_engine/model/embedding_model/bge-m3-model \
  --tokenizer_path ./embedding_engine/model/embedding_model/bge-m3-tokenizer \
  --output_path ./embedding_models/model_dynamic/bge_m3_embedding_dynamic.onnx \
  --max_length 256

參數與腳本說明 (點擊展開)

參數	預設值	說明
`--model_path`	`./embedding_engine/model/embedding_model/bge-m3-model`	模型路徑
`--tokenizer_path`	`./embedding_engine/model/embedding_model/bge-m3-tokenizer`	Tokenizer 路徑
`--output_path`	`./embedding_models/model_dynamic/bge_m3_embedding_dynamic.onnx`	輸出 ONNX 路徑
`--max_length`	`256`	最大序列長度

腳本說明：

使用 AutoModel 載入模型，並透過 EmbeddingWrapper 擷取 CLS token（last_hidden_state[:, 0, :]）作為輸出
輸入：input_ids、attention_mask（int32）；輸出：embeddings
動態軸：batch_size（第 0 維），ONNX opset：17

Reranker 模型（BGE-Reranker-Large）

python trt_convert/rerank2onnx.py \
  --model_path ./reranking_model/bge-reranker-large-model \
  --tokenizer_path ./reranking_model/bge-reranker-large-tokenizer \
  --output_path ./model_dynamic/bge_reranker_large_dynamic.onnx \
  --max_length 256

參數與腳本說明 (點擊展開)

參數	預設值	說明
`--model_path`	`./reranking_model/bge-reranker-large-model`	模型路徑
`--tokenizer_path`	`./reranking_model/bge-reranker-large-tokenizer`	Tokenizer 路徑
`--output_path`	`./model_dynamic/bge_reranker_large_dynamic.onnx`	輸出 ONNX 路徑
`--max_length`	`256`	最大序列長度

腳本說明：

使用 AutoModelForSequenceClassification 載入模型，RerankerWrapper 輸出 logits.squeeze(-1) 作為相關性分數
輸入：input_ids、attention_mask（int32）；輸出：scores
動態軸：batch_size（第 0 維），ONNX opset：17

NLI 模型（XLM-RoBERTa-Large-XNLI）

python trt_convert/nli2onnx.py \
  --model_path joeddav/xlm-roberta-large-xnli \
  --tokenizer_path joeddav/xlm-roberta-large-xnli \
  --output_path ./model_dynamic_bs/nli_model_dynamic_bs.onnx \
  --max_length 256

參數與腳本說明 (點擊展開)

參數	預設值	說明
`--model_path`	`joeddav/xlm-roberta-large-xnli`	模型路徑（支援 HuggingFace Hub）
`--tokenizer_path`	`joeddav/xlm-roberta-large-xnli`	Tokenizer 路徑
`--output_path`	`./model_dynamic_bs/nli_model_dynamic_bs.onnx`	輸出 ONNX 路徑
`--max_length`	`256`	最大序列長度

腳本說明：

使用 AutoModelForSequenceClassification 載入模型，直接輸出 logits（3 個 NLI 類別分數）
輸入：input_ids、attention_mask（int32）；輸出：logits
動態軸：batch_size（第 0 維），ONNX opset：17

步驟二：ONNX 轉換為 TensorRT

需在 TensorRT Docker 容器內執行（可參考 docker/docker-compose.yaml），使用 trtexec 工具進行轉換。

靜態 Batch Size

trtexec \
  --onnx=./model/nli_model_dynamic_bs.onnx \
  --saveEngine=nli_model_bs8.trt \
  --fp16

動態 Batch Size（建議）

動態 batch size 允許模型在推論時接受不同大小的批次。以下以各模型為例：

Embedding 模型

trtexec \
  --onnx=./embedding_models/bge_m3_embedding_dynamic.onnx \
  --saveEngine=bge_m3_model_dynamic_bs.trt \
  --fp16 \
  --minShapes=input_ids:1x256,attention_mask:1x256 \
  --optShapes=input_ids:8x256,attention_mask:8x256 \
  --maxShapes=input_ids:32x256,attention_mask:32x256

Reranker 模型

trtexec \
  --onnx=./reranker_models/bge_reranker_large_dynamic.onnx \
  --saveEngine=bge_reranker_large_dynamic_bs.trt \
  --fp16 \
  --minShapes=input_ids:1x256,attention_mask:1x256 \
  --optShapes=input_ids:8x256,attention_mask:8x256 \
  --maxShapes=input_ids:32x256,attention_mask:32x256

NLI 模型

trtexec \
  --onnx=./nli_models/nli_model_dynamic_bs.onnx \
  --saveEngine=nli_model_dynamic_bs.trt \
  --fp16 \
  --minShapes=input_ids:1x256,attention_mask:1x256 \
  --optShapes=input_ids:8x256,attention_mask:8x256 \
  --maxShapes=input_ids:32x256,attention_mask:32x256

trtexec 參數說明

參數	說明
`--onnx`	輸入的 ONNX 模型路徑
`--saveEngine`	輸出的 TensorRT 引擎路徑
`--fp16`	啟用 FP16 精度以加速推論並減少記憶體使用
`--minShapes`	動態形狀的最小輸入尺寸 `tensor_name:dim0xdim1`
`--optShapes`	動態形狀的最佳化輸入尺寸（影響 TRT 優化重點）
`--maxShapes`	動態形狀的最大輸入尺寸

轉換完成後，將 .trt 檔案路徑填入 configs/config.yaml 的 model_path 欄位即可。

啟動 Server 配置文件

編輯 configs/config.yaml 來配置您的模型：

nli_models:
  xlm-roberta-large-xnli:
    model_name: "xlm-roberta-large-xnli"
    model_path: "./model/nlimodels/trtmodels/nli_model_dynamic_bs.trt"
    tokenizer_path: "joeddav/xlm-roberta-large-xnli"
    reuse_dynamic_buffer: true
    cuda_graph_list:
      - 1
      - 3
      - 5

embedding_models:
  bge-m3:
    model_name: "bge-m3"
    model_path: "./model/embedding_models/trt_models/bge_m3_model_dynamic_bs.trt"
    tokenizer_path: "./model/embedding_models/bge-m3-tokenizer"
    reuse_dynamic_buffer: true
    cuda_graph_list:
      - 1
      - 3
      - 5

reranking_models:
  bge-reranker-large:
    model_name: "bge-reranker-large"
    model_path: "./model/reranker_models/trt_models/bge_reranker_large_dynamic_bs.trt"
    tokenizer_path: "./model/reranker_models/bge-reranker-large-tokenizer"
    reuse_dynamic_buffer: true
    cuda_graph_list:
      - 1
      - 3
      - 5

配置說明

model_name: 模型識別名稱
model_path: TensorRT 模型文件路徑
tokenizer_path: Tokenizer 路徑（可以是本地路徑或 Hugging Face 模型名稱）
reuse_dynamic_buffer: 動態 batch size 是否初始化預先分配 buffer，推輪時就不須每次動態分配
cuda_graph_list: 根據指定 batch size，預先生成一串固定的 GPU kernel 調用，減少 CPU→GPU kernel launch 的 overhead

服務啟動

使用腳本啟動

chmod +x start_tensorrt_server.sh
./start_tensorrt_server.sh

服務啟動後，API 將在 http://{ip}:{port} 上提供服務。

API 使用說明

服務提供兩種 API 格式：原生 API 和 OpenAI 相容 API。

查看可用模型

curl http://localhost:8887/models

回應：

{
  "models": {
    "embedding_models": ["bge-m3"],
    "reranking_models": ["bge-reranker-large"],
    "nli_models": ["xlm-roberta-large-xnl"]
  }
}

原生 API 格式

使用方式說明 (點擊展開)

1. Embedding 推論

import requests

url = "http://localhost:8887/infer/bge-m3"
payload = {"documents": ["這是一個測試文本", "另一個測試文本"]}
response = requests.post(url, json=payload)

print(response.json())
# 輸出:
# {
#     "embeddings": [[0.1, 0.2, ...], [0.3, 0.4, ...]],
#     "elapsed_ms": 15.2
# }

2. Reranker 推論

import requests

url = "http://localhost:8887/infer/bge-reranker-large"
payload = {
    "query": "機器學習的理論很重要",
    "documents": [
        "理論對於理解機器學習非常重要",
        "實作經驗在機器學習中也很關鍵"
    ]
}
response = requests.post(url, json=payload)

print(response.json())
# 輸出:
# {
#     "scores": [9.5078125, 7.2421875],
#     "elapsed_ms": 5.78
# }

3. NLI (自然語言推理) 推論

import requests

url = "http://localhost:8887/infer/xlm-roberta-large-xnli"
payload = {
    "premises": ["今天天氣很好", "貓是動物"],
    "hypotheses": ["今天是晴天", "狗是動物"]
}
response = requests.post(url, json=payload)

print(response.json())
# 輸出:
# {
#     "predictions": ["entailment", "neutral"],
#     "logits": [[2.18, -1.38, -0.72], [1.02, -0.42, -0.65]],
#     "elapsed_ms": 12.45
# }

OpenAI 相容 API

注意: 僅支援 Embedding 和 Reranker 模型，NLI 模型不支援此格式。

使用方式說明 (點擊展開)

1. Embedding API

from openai import OpenAI

client = OpenAI(
    api_key="EMPTY",  # 任意值即可
    base_url="http://localhost:8887/v1"
)

text = "這是一個測試文本"
response = client.embeddings.create(
    input=[text],
    model="bge-m3"
)

print(response.data[0].embedding)

2. Reranker API

from openai import OpenAI

client = OpenAI(
    api_key="EMPTY",
    base_url="http://localhost:8887/v1"
)

documents = [
    "機器學習最好通過專案來學習",
    "理論對於理解機器學習很重要"
]

response = client.embeddings.create(
    model="bge-reranker-large",
    input=documents,
    extra_body={"query": "理論對於理解機器學習很重要"}
)

# 取得重排序分數
scores = [data.embedding for data in response.data]
print(scores)

請求參數說明

Embedding 請求

documents (required): 字串或字串列表，要編碼的文本
model (optional): 模型名稱，在 OpenAI API 中為必需

Reranker 請求

query (required): 查詢字串
documents (required): 候選文件列表
model (optional): 模型名稱，在 OpenAI API 中為必需

NLI 請求

premises (required): 前提句子列表
hypotheses (required): 假設句子列表

回應格式說明

Embedding 回應

{
  "embeddings": [[0.1, 0.2, ...]], // 向量列表
  "elapsed_ms": 15.2               // 推論時間（毫秒）
}

Reranker 回應

{
  "scores": [9.5078125],           // 相關性分數列表
  "elapsed_ms": 5.78               // 推論時間（毫秒）
}

NLI 回應

{
  "predictions": ["entailment"],   // 預測標籤列表
  "logits": [[2.18, -1.38, -0.72]], // 原始分數
  "elapsed_ms": 12.45              // 推論時間（毫秒）
}

NLI 標籤說明

entailment: 蘊含（前提支持假設）
neutral: 中性（前提與假設無關）
contradiction: 矛盾（前提與假設衝突）

性能基準測試

以下為在 NVIDIA A100 GPU 上的測試結果：

Embedding 推論性能

測試配置:

模型: BGE-M3
批次大小: 1-64
推論比較對象: torch v.s TensorRT
推論 Server 對象: 串行推論 v.s 動態組批

Reranker 推論性能

測試配置:

模型: BGE-Reranker-Large
批次大小: 1-64
推論比較對象: torch v.s TensorRT
推論 Server 對象: 串行推論 v.s 動態組批

NLI 推論性能

測試配置:

模型: XLM-RoBERTa-Large-XNLI
批次大小: 1-64
推論比較對象: torch v.s TensorRT
推論 Server 對象: 串行推論 v.s 動態組批

架構設計

系統架構圖

sequenceDiagram
    participant Client as 客戶端
    participant API as FastAPI 服務
    participant Worker as Worker 處理器
    participant TRT as TensorRT 推論器
    participant GPU as GPU (CUDA + TensorRT)

    Client->>API: HTTP 請求 (/infer, /v1/embeddings)
    API->>Worker: 動態佇列 (payload, Future)
    Worker->>Worker: 收集批次 (max_batch=32, max_wait=10ms)
    Worker->>TRT: 呼叫 model.infer(all_docs)

    TRT->>TRT: Tokenizer → input_ids, attention_mask
    TRT->>TRT: 轉換資料型別 → engine dtype
    TRT->>GPU: H2D (memcpy_htod_async)
    Note right of GPU: GPU 緩衝區準備就緒
    TRT->>GPU: set_input_shape / set_tensor_address
    TRT->>GPU: execute_async_v3()
    GPU-->>TRT: GPU 推論完成
    Note over GPU: Engine 在 GPU 上執行推論
    TRT->>GPU: D2H (memcpy_dtoh_async)
    TRT->>TRT: 重新整形，截斷至原始長度
    TRT-->>Worker: 預測結果 / logits

    Worker->>Worker: 分割批次結果
    Worker-->>API: Future.set_result()
    API-->>Client: 回傳 JSON 結果

License

This project is licensed under the MIT License.

FilesExpand file tree

README_zh-CN.md

Latest commit

History

README_zh-CN.md

File metadata and controls

TensorRT 模型推論服務器

功能特色

支援的模型

安裝與設定

1. 克隆專案

2. 安裝依賴

轉換模型

步驟一：轉換為 ONNX

Embedding 模型（BGE-M3）

Reranker 模型（BGE-Reranker-Large）

NLI 模型（XLM-RoBERTa-Large-XNLI）

步驟二：ONNX 轉換為 TensorRT

靜態 Batch Size

動態 Batch Size（建議）

trtexec 參數說明

啟動 Server 配置文件

配置說明

服務啟動

使用腳本啟動

API 使用說明

查看可用模型

原生 API 格式

1. Embedding 推論

2. Reranker 推論

3. NLI (自然語言推理) 推論

OpenAI 相容 API

1. Embedding API

2. Reranker API

請求參數說明

Embedding 請求

Reranker 請求

NLI 請求

回應格式說明

Embedding 回應

Reranker 回應

NLI 回應

NLI 標籤說明

性能基準測試

Embedding 推論性能

Reranker 推論性能

NLI 推論性能

架構設計

系統架構圖

License