sentence-transformers GPU 推理 3.87 倍提速指南：三步跑通 ONNX 后端

【免费下载链接】sentence-transformersState-of-the-Art Embeddings, Retrieval, and Reranking项目地址: https://gitcode.com/gh_mirrors/se/sentence-transformers

还在用 PyTorch FP32 在 GPU 上批量嵌入百万级句子？本文用三步把 sentence-transformers 搬上 ONNX 后端，GPU 中位吞吐量最高提到 ⚡3.87 倍。先看数据。

先说结果

下表来自官方仓库的基准测试：每行一种推理配置，提速均相对默认的 PyTorch FP32 基线取中位数。

测试硬件是一块 RTX 3090。模型覆盖 MiniLM（22.7M 参数）到 BGE-M3（567M 参数）共 4 个。负载分短句子、问答答案、长评三类。还有一个让人放心的数字：这些配置的平均任务质量比都在 1 附近，也就是说提速没有可测量的质量损失。

原理一分钟：为什么能快

PyTorch 默认是动态图，每次前向都在运行时决定执行计划。批里的短文本会被 pad 到最长，大量算力浪费在填充位上。

提速来自四个杠杆：

• 半精度。FP16（半精度浮点，数值宽度减半）和 BF16 让 GPU 上的 tensor core（专门跑低精度矩阵运算的硬件单元）满频工作，中位加速直接到 2.92x。
• Flash Attention 2。一种为 GPU 显存访问优化的注意力实现，配合"去填充"（unpadding，把 pad 位剔掉、只算真实 token），整批拼成一条连续序列，中位加速爬到 3.87x。
• ONNX（开放的可移植模型计算图格式）后端。ONNX Runtime 用静态图接管：导出后计算图固定，算子融合和内存布局提前规划好。
• 算力外包。执行提供方（execution provider）是 ONNX Runtime 里真正跑算子的计算后端，可以是 CUDA（NVIDIA 的 GPU 计算平台），也可以走 TensorRT（NVIDIA 的 GPU 推理引擎）。

图里可以看到：批内文本长度越不均匀，去填充收益越大，长短混合的负载提升最明显。

三步把 ONNX 后端跑通

以下步骤统一用 all-MiniLM-L6-v2 举例。

第一步｜装好 ONNX 运行时

目标：装上能执行 ONNX 模型的运行时。

# GPU 加速推理装 onnx-gpu 附加组件 pip install sentence-transformers[onnx-gpu]

最常见的坑：只装了裸的pip install sentence-transformers没带附加组件，再写backend="onnx"加载会直接抛 ImportError，看起来像模型问题，其实是缺依赖。完整依赖清单见 installation.md。

第二步｜一行命令完成模型转换

目标：不写任何转换脚本，从预训练模型拿到 ONNX 产物。

from sentence_transformers import SentenceTransformer # 模型源里没有 ONNX 文件时，首次加载自动导出 model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2", backend="onnx") # 把产物存到本地，下次启动不用重付导出成本 model.save_pretrained("./all-minilm-onnx")

坑：导出本身要花几十秒，不存盘的话每次进程启动都重付一遍。存一次、之后加载本地目录，才是上线姿势。

第三步｜指定 GPU 提供方跑推理

目标：加载本地 ONNX 产物，把计算指到 GPU 上，跑通一次编码验证输出。

model = SentenceTransformer( "./all-minilm-onnx", backend="onnx", model_kwargs={"provider": "CUDAExecutionProvider"}, # 指定 GPU 执行后端 ) embeddings = model.encode(["GPU 推理必备", "sentence-transformers ONNX 后端"])

provider选的是 ONNX Runtime 的执行提供方。不指定时，框架自动取环境里优先级最高的提供方，顺序是 TensorRT > CUDA > CPU，逻辑见 backend/load.py。

坑：回退到低优先级提供方不会报错，只会加速比看着差。部署前先用onnxruntime.get_available_providers()确认你期望的提供方真的在列表里。

真上线前的坑

1. 加载不慢，加速比却接近 0。原因：模型源里已带一份未优化的 ONNX 文件，被自动加载了。解法：先跑一次export_optimized_onnx_model推到 O3/O4，加载时用file_name指定优化后的产物。
2. 输出维度和数值跟 PyTorch 版对不上。原因：在框架外直接裸用 ONNX 文件，而导出只转了 Transformer 部分，输出的是 token 嵌入序列。解法：走.encode()让框架补池化和归一化，或者自己实现池化。
3. 小批量单请求延迟改善不明显。原因：表中数字来自最佳批量下的吞吐中位数，小模型上分词和 Python 开销占大头。解法：加大批量，或改走 PyTorch 路径用model.compile(mode="reduce-overhead")加固定长度 padding。
4. 提供方不是你以为的那个。原因：环境缺对应运行时，ONNX Runtime 静默选了低优先级提供方。🔧 解法：部署前用onnxruntime.get_available_providers()核对可用提供方列表。
5. 和生成式 LLM 同卡，显存飙升、两边都变慢。原因：两个模型争抢显存和算力。解法：延迟敏感的本地部署把小嵌入模型挪到 CPU（device="cpu"），或者分卡部署。

如果还不够快

• ONNX 图优化 + 量化。export_optimized_onnx_model支持 O1–O4 四档，O4 在 GPU 上额外启用 fp16 混合精度；CPU 侧可用export_dynamic_quantized_onnx_model转 int8。实现都在 backend/optimize.py。
• 不切 ONNX 的 PyTorch 路径。model.compile()加 bf16 加 Flash Attention 2，CUDA 图模式需要固定长度输入，在processing_kwargs传padding="max_length"。完整参数说明在 efficiency.rst。
• 存储和检索成本。瓶颈在索引体积时，对嵌入向量做后置量化（把向量压成更低精度再入库），见sentence_transformers/util/quantization.py。

如果你的硬件或负载和上面不同，先拿自己的数据跑一遍仓库自带的基准脚本再下结论：evaluation_inference_speed.py。

把上面几步配置到位，多数嵌入服务的 GPU 利用率和延迟都能落到一个健康区间。碰到提供方回退、质量回退这类意外，把日志带上直接去项目 issue 区提问；自己组合出更优配置，也可以提 PR 回馈给仓库。

【免费下载链接】sentence-transformersState-of-the-Art Embeddings, Retrieval, and Reranking项目地址: https://gitcode.com/gh_mirrors/se/sentence-transformers