加速推理

Sentence Transformers 支持 3 种后端，用于使用 Cross Encoder 模型执行推理，每种后端都有自己的优化，可加速推理。

Cross Encoder 的默认后端。灵活高效的模型加速器。模型优化，主要针对英特尔硬件。不同后端的基准测试。用于导出、优化和量化模型的 GUI。

PyTorch

PyTorch 后端是 Cross Encoder 的默认后端。如果您不指定设备，它将使用“cuda”、“mps”和“cpu”中最强的可用选项。其默认用法如下所示：

from sentence_transformers import CrossEncoder

model = CrossEncoder("cross-encoder/ms-marco-MiniLM-L6-v2")

query = "Which planet is known as the Red Planet?"
passages = [
   "Venus is often called Earth's twin because of its similar size and proximity.",
   "Mars, known for its reddish appearance, is often referred to as the Red Planet.",
   "Jupiter, the largest planet in our solar system, has a prominent red spot.",
   "Saturn, famous for its rings, is sometimes mistaken for the Red Planet."
]

scores = model.predict([(query, passage) for passage in passages])
print(scores)

如果您正在使用 GPU，则可以使用以下选项来加速推理：

float16 (fp16)

Float32 (fp32，全精度) 是 torch 中默认的浮点格式，而 float16 (fp16，半精度) 是一种精度降低的浮点格式，可以在 GPU 上加速推理，同时模型精度损失极小。要使用它，您可以在初始化期间指定 torch_dtype，或者在已初始化的模型上调用 model.half()。

from sentence_transformers import CrossEncoder

model = CrossEncoder("cross-encoder/ms-marco-MiniLM-L6-v2", model_kwargs={"torch_dtype": "float16"})
# or: model.half()

query = "Which planet is known as the Red Planet?"
passages = [
   "Venus is often called Earth's twin because of its similar size and proximity.",
   "Mars, known for its reddish appearance, is often referred to as the Red Planet.",
   "Jupiter, the largest planet in our solar system, has a prominent red spot.",
   "Saturn, famous for its rings, is sometimes mistaken for the Red Planet."
]

scores = model.predict([(query, passage) for passage in passages])
print(scores)

bfloat16 (bf16)

Bfloat16 (bf16) 类似于 fp16，但保留了更多 fp32 的原始精度。要使用它，您可以在初始化期间指定 torch_dtype，或者在已初始化的模型上调用 model.bfloat16()。

from sentence_transformers import CrossEncoder

model = CrossEncoder("cross-encoder/ms-marco-MiniLM-L6-v2", model_kwargs={"torch_dtype": "bfloat16"})
# or: model.bfloat16()

query = "Which planet is known as the Red Planet?"
passages = [
   "Venus is often called Earth's twin because of its similar size and proximity.",
   "Mars, known for its reddish appearance, is often referred to as the Red Planet.",
   "Jupiter, the largest planet in our solar system, has a prominent red spot.",
   "Saturn, famous for its rings, is sometimes mistaken for the Red Planet."
]

scores = model.predict([(query, passage) for passage in passages])
print(scores)

ONNX

ONNX 可用于通过将模型转换为 ONNX 格式并使用 ONNX Runtime 运行模型来加速推理。要使用 ONNX 后端，您必须安装 Sentence Transformers 并分别使用 onnx 或 onnx-gpu 额外依赖项以进行 CPU 或 GPU 加速。

pip install sentence-transformers[onnx-gpu]
# or
pip install sentence-transformers[onnx]

要将模型转换为 ONNX 格式，您可以使用以下代码：

from sentence_transformers import CrossEncoder

model = CrossEncoder("cross-encoder/ms-marco-MiniLM-L6-v2", backend="onnx")

query = "Which planet is known as the Red Planet?"
passages = [
   "Venus is often called Earth's twin because of its similar size and proximity.",
   "Mars, known for its reddish appearance, is often referred to as the Red Planet.",
   "Jupiter, the largest planet in our solar system, has a prominent red spot.",
   "Saturn, famous for its rings, is sometimes mistaken for the Red Planet."
]

scores = model.predict([(query, passage) for passage in passages])
print(scores)

如果模型路径或仓库中已包含 ONNX 格式的模型，Sentence Transformers 将自动使用它。否则，它会将模型转换为 ONNX 格式。

注意

如果您希望在 Sentence Transformers 之外使用 ONNX 模型，您可能需要应用您选择的激活函数（例如 Sigmoid）才能获得与 Sentence Transformers 中 Cross Encoder 相同的结果。

所有通过 model_kwargs 传递的关键字参数都将传递给 ORTModelForSequenceClassification.from_pretrained。一些值得注意的参数包括：

provider：用于加载模型的 ONNX Runtime 提供者，例如 "CPUExecutionProvider"。有关可能的提供者，请参阅 https://runtime.onnx.org.cn/docs/execution-providers/。如果未指定，将使用最强的提供者（例如 "CUDAExecutionProvider"）。
file_name：要加载的 ONNX 文件的名称。如果未指定，将默认为 "model.onnx"，否则为 "onnx/model.onnx"。此参数对于指定优化或量化模型很有用。
export：一个布尔标志，指定模型是否将被导出。如果未提供，如果模型仓库或目录中尚未包含 ONNX 模型，则 export 将设置为 True。

提示

强烈建议保存导出的模型，以避免每次运行代码时都重新导出。如果您的模型是本地的，您可以通过调用 model.save_pretrained() 来实现。

model = CrossEncoder("path/to/my/model", backend="onnx")
model.save_pretrained("path/to/my/model")

或者如果您的模型来自 Hugging Face Hub，则使用 model.push_to_hub()。

model = CrossEncoder("Alibaba-NLP/gte-reranker-modernbert-base", backend="onnx")
model.push_to_hub("Alibaba-NLP/gte-reranker-modernbert-base", create_pr=True)

优化 ONNX 模型

ONNX 模型可以使用 Optimum 进行优化，从而在 CPU 和 GPU 上都能实现加速。为此，您可以使用 export_optimized_onnx_model() 函数，该函数将优化后的模型保存在您指定的目录或模型仓库中。它期望：

model：一个使用 ONNX 后端加载的 Sentence Transformer 或 Cross Encoder 模型。
optimization_config："O1"、"O2"、"O3" 或 "O4"，表示来自 AutoOptimizationConfig 的优化级别，或一个 OptimizationConfig 实例。
model_name_or_path：用于保存优化模型文件的路径，如果您想将其推送到 Hugging Face Hub，则为仓库名称。
push_to_hub：（可选）一个布尔值，用于将优化模型推送到 Hugging Face Hub。
create_pr：（可选）一个布尔值，用于在推送到 Hugging Face Hub 时创建拉取请求。当您对仓库没有写入权限时很有用。
file_suffix：（可选）保存模型时要附加到模型名称的字符串。如果未指定，将使用优化级别名称字符串，如果优化配置不是字符串优化级别，则仅使用 "optimized"。

有关使用优化级别 3（基本和扩展的通用优化、transformers 特定的融合、快速 Gelu 近似）导出模型的示例，请参见此处。

Hugging Face Hub 模型

只优化一次

from sentence_transformers import CrossEncoder, export_optimized_onnx_model

model = CrossEncoder("cross-encoder/ms-marco-MiniLM-L6-v2", backend="onnx")
export_optimized_onnx_model(
    model,
    "O3",
    "cross-encoder/ms-marco-MiniLM-L6-v2",
    push_to_hub=True,
    create_pr=True,
)

在拉取请求合并之前

from sentence_transformers import CrossEncoder

pull_request_nr = 2 # TODO: Update this to the number of your pull request
model = CrossEncoder(
    "cross-encoder/ms-marco-MiniLM-L6-v2",
    backend="onnx",
    model_kwargs={"file_name": "onnx/model_O3.onnx"},
    revision=f"refs/pr/{pull_request_nr}"
)

一旦拉取请求合并

from sentence_transformers import CrossEncoder

model = CrossEncoder(
    "cross-encoder/ms-marco-MiniLM-L6-v2",
    backend="onnx",
    model_kwargs={"file_name": "onnx/model_O3.onnx"},
)

本地模型

只优化一次

from sentence_transformers import CrossEncoder, export_optimized_onnx_model

model = CrossEncoder("path/to/my/mpnet-legal-finetuned", backend="onnx")
export_optimized_onnx_model(model, "O3", "path/to/my/mpnet-legal-finetuned")

优化后

from sentence_transformers import CrossEncoder

model = CrossEncoder(
    "path/to/my/mpnet-legal-finetuned",
    backend="onnx",
    model_kwargs={"file_name": "onnx/model_O3.onnx"},
)

量化 ONNX 模型

ONNX 模型可以使用 Optimum 量化为 int8 精度，从而在 CPU 上实现更快的推理。为此，您可以使用 export_dynamic_quantized_onnx_model() 函数，该函数将量化后的模型保存在您指定的目录或模型仓库中。动态量化与静态量化不同，不需要校准数据集。它期望：

model：一个使用 ONNX 后端加载的 Sentence Transformer 或 Cross Encoder 模型。
quantization_config："arm64"、"avx2"、"avx512" 或 "avx512_vnni"，表示来自 AutoQuantizationConfig 的量化配置，或一个 QuantizationConfig 实例。
model_name_or_path：用于保存量化模型文件的路径，如果您想将其推送到 Hugging Face Hub，则为仓库名称。
push_to_hub：（可选）一个布尔值，用于将量化模型推送到 Hugging Face Hub。
create_pr：（可选）一个布尔值，用于在推送到 Hugging Face Hub 时创建拉取请求。当您对仓库没有写入权限时很有用。
file_suffix：（可选）保存模型时要附加到模型名称的字符串。如果未指定，将使用 "qint8_quantized"。

在我的 CPU 上，每个默认量化配置（"arm64"、"avx2"、"avx512"、"avx512_vnni"）都带来了大致相当的加速效果。

有关使用 avx512_vnni 将模型量化为 int8 的示例，请参见此处。

Hugging Face Hub 模型

只量化一次

from sentence_transformers import CrossEncoder, export_dynamic_quantized_onnx_model

model = CrossEncoder("cross-encoder/ms-marco-MiniLM-L6-v2", backend="onnx")
export_dynamic_quantized_onnx_model(
    model,
    "avx512_vnni",
    "sentence-transformers/cross-encoder/ms-marco-MiniLM-L6-v2",
    push_to_hub=True,
    create_pr=True,
)

在拉取请求合并之前

from sentence_transformers import CrossEncoder

pull_request_nr = 2 # TODO: Update this to the number of your pull request
model = CrossEncoder(
    "cross-encoder/ms-marco-MiniLM-L6-v2",
    backend="onnx",
    model_kwargs={"file_name": "onnx/model_qint8_avx512_vnni.onnx"},
    revision=f"refs/pr/{pull_request_nr}",
)

一旦拉取请求合并

from sentence_transformers import CrossEncoder

model = CrossEncoder(
    "cross-encoder/ms-marco-MiniLM-L6-v2",
    backend="onnx",
    model_kwargs={"file_name": "onnx/model_qint8_avx512_vnni.onnx"},
)

本地模型

只量化一次

from sentence_transformers import CrossEncoder, export_dynamic_quantized_onnx_model

model = CrossEncoder("path/to/my/mpnet-legal-finetuned", backend="onnx")
export_dynamic_quantized_onnx_model(model, "O3", "path/to/my/mpnet-legal-finetuned")

量化后

from sentence_transformers import CrossEncoder

model = CrossEncoder(
    "path/to/my/mpnet-legal-finetuned",
    backend="onnx",
    model_kwargs={"file_name": "onnx/model_qint8_avx512_vnni.onnx"},
)

OpenVINO

OpenVINO 允许通过将模型导出为 OpenVINO 格式来在 CPU 上加速推理。要使用 OpenVINO 后端，您必须安装带有 openvino 额外依赖项的 Sentence Transformers。

pip install sentence-transformers[openvino]

要将模型转换为 OpenVINO 格式，您可以使用以下代码：

from sentence_transformers import CrossEncoder

model = CrossEncoder("cross-encoder/ms-marco-MiniLM-L6-v2", backend="openvino")

query = "Which planet is known as the Red Planet?"
passages = [
   "Venus is often called Earth's twin because of its similar size and proximity.",
   "Mars, known for its reddish appearance, is often referred to as the Red Planet.",
   "Jupiter, the largest planet in our solar system, has a prominent red spot.",
   "Saturn, famous for its rings, is sometimes mistaken for the Red Planet."
]

scores = model.predict([(query, passage) for passage in passages])
print(scores)

如果模型路径或仓库中已包含 OpenVINO 格式的模型，Sentence Transformers 将自动使用它。否则，它会将模型转换为 OpenVINO 格式。

注意

如果您希望在 Sentence Transformers 之外使用 OpenVINO 模型，您可能需要应用您选择的激活函数（例如 Sigmoid）才能获得与 Sentence Transformers 中 Cross Encoder 相同的结果。

所有通过 model_kwargs 传递的关键字参数都将传递给 OVBaseModel.from_pretrained()。一些值得注意的参数包括：

file_name：要加载的 ONNX 文件的名称。如果未指定，将默认为 "openvino_model.xml"，否则为 "openvino/openvino_model.xml"。此参数对于指定优化或量化模型很有用。
export：一个布尔标志，指定模型是否将被导出。如果未提供，如果模型仓库或目录中尚未包含 OpenVINO 模型，则 export 将设置为 True。

提示

强烈建议保存导出的模型，以避免每次运行代码时都重新导出。如果您的模型是本地的，您可以通过调用 model.save_pretrained() 来实现。

model = CrossEncoder("path/to/my/model", backend="openvino")
model.save_pretrained("path/to/my/model")

或者如果您的模型来自 Hugging Face Hub，则使用 model.push_to_hub()。

model = CrossEncoder("Alibaba-NLP/gte-reranker-modernbert-base", backend="openvino")
model.push_to_hub("Alibaba-NLP/gte-reranker-modernbert-base", create_pr=True)

量化 OpenVINO 模型

OpenVINO 模型可以使用 Optimum Intel 量化为 int8 精度以加速推理。为此，您可以使用 export_static_quantized_openvino_model() 函数，该函数将量化后的模型保存在您指定的目录或模型仓库中。训练后静态量化期望：

model：一个使用 OpenVINO 后端加载的 Sentence Transformer 或 Cross Encoder 模型。
quantization_config：（可选）量化配置。此参数接受以下任一值：None 表示默认的 8 位量化，一个表示量化配置的字典，或一个 OVQuantizationConfig 实例。
model_name_or_path：用于保存量化模型文件的路径，如果您想将其推送到 Hugging Face Hub，则为仓库名称。
dataset_name：（可选）用于校准的数据集的名称。如果未指定，默认为 glue 数据集中的 sst2 子集。
dataset_config_name：（可选）要加载的数据集的特定配置。
dataset_split：（可选）要加载的数据集分割（例如，“train”、“test”）。
column_name：（可选）数据集中用于校准的列名。
push_to_hub：（可选）一个布尔值，用于将量化模型推送到 Hugging Face Hub。
create_pr：（可选）一个布尔值，用于在推送到 Hugging Face Hub 时创建拉取请求。当您对仓库没有写入权限时很有用。
file_suffix：（可选）保存模型时要附加到模型名称的字符串。如果未指定，将使用 "qint8_quantized"。

有关使用静态量化将模型量化为 int8 的示例，请参见此处。

Hugging Face Hub 模型

只量化一次

from sentence_transformers import CrossEncoder, export_static_quantized_openvino_model

model = CrossEncoder("cross-encoder/ms-marco-MiniLM-L6-v2", backend="openvino")
export_static_quantized_openvino_model(
    model,
    quantization_config=None,
    model_name_or_path="cross-encoder/ms-marco-MiniLM-L6-v2",
    push_to_hub=True,
    create_pr=True,
)

在拉取请求合并之前

from sentence_transformers import CrossEncoder

pull_request_nr = 2 # TODO: Update this to the number of your pull request
model = CrossEncoder(
    "cross-encoder/ms-marco-MiniLM-L6-v2",
    backend="openvino",
    model_kwargs={"file_name": "openvino/openvino_model_qint8_quantized.xml"},
    revision=f"refs/pr/{pull_request_nr}"
)

一旦拉取请求合并

from sentence_transformers import CrossEncoder

model = CrossEncoder(
    "cross-encoder/ms-marco-MiniLM-L6-v2",
    backend="openvino",
    model_kwargs={"file_name": "openvino/openvino_model_qint8_quantized.xml"},
)

本地模型

只量化一次

from sentence_transformers import CrossEncoder, export_static_quantized_openvino_model
from optimum.intel import OVQuantizationConfig

model = CrossEncoder("path/to/my/mpnet-legal-finetuned", backend="openvino")
quantization_config = OVQuantizationConfig()
export_static_quantized_openvino_model(model, quantization_config, "path/to/my/mpnet-legal-finetuned")

量化后

from sentence_transformers import CrossEncoder

model = CrossEncoder(
    "path/to/my/mpnet-legal-finetuned",
    backend="openvino",
    model_kwargs={"file_name": "openvino/openvino_model_qint8_quantized.xml"},
)

基准测试

以下图像显示了不同后端在 GPU 和 CPU 上的基准测试结果。结果是根据 4 种不同大小的模型、3 个数据集和多种批次大小的平均值。

展开基准测试详情

加速比

硬件：RTX 3090 GPU，i7-17300K CPU
数据集：GPU 测试使用 2000 个样本，CPU 测试使用 1000 个样本。
- sentence-transformers/stsb：sentence1 和 sentence2 列作为配对，平均分别有 38.94 ± 13.97 和 38.96 ± 14.05 个字符。
- sentence-transformers/natural-questions：query 和 answer 列作为配对，平均分别有 46.99 ± 10.98 和 619.63 ± 345.30 个字符。
- stanfordnlp/imdb：从 text 列中使用了两种变体：前 100 个字符（100.00 ± 0.00 个字符）和每个样本重复 4 次（16804.25 ± 10178.26 个字符）。
模型
- cross-encoder/ms-marco-MiniLM-L6-v2：22.7M 参数；批次大小为 16、32、64、128 和 256。
- BAAI/bge-reranker-base：278M 参数；批次大小为 16、32、64 和 128。
- mixedbread-ai/mxbai-rerank-large-v1：435M 参数；批次大小为 8、16、32 和 64。GPU 测试也使用了 128 和 256。
- BAAI/bge-reranker-v2-m3：568M 参数；批次大小为 2、4。GPU 测试也使用了 8、16 和 32。

性能比：使用了相同的模型和硬件。我们将性能与使用 fp32 的 PyTorch（即默认后端和精度）的性能进行比较。

评估
- 信息检索：基于 MS MARCO 和 NQ 子集（来自 NanoBEIR 数据集集合）的余弦相似度计算的 NDCG@10，通过 CrossEncoderNanoBEIREvaluator 计算。

后端
- torch-fp32：PyTorch 使用 float32 精度（默认）。
- torch-fp16：PyTorch 使用 float16 精度，通过 model_kwargs={"torch_dtype": "float16"}。
- torch-bf16：PyTorch 使用 bfloat16 精度，通过 model_kwargs={"torch_dtype": "bfloat16"}。
- onnx：ONNX 使用 float32 精度，通过 backend="onnx"。
- onnx-O1：ONNX 使用 float32 精度和 O1 优化，通过 export_optimized_onnx_model(..., "O1", ...) 和 backend="onnx"。
- onnx-O2：ONNX 使用 float32 精度和 O2 优化，通过 export_optimized_onnx_model(..., "O2", ...) 和 backend="onnx"。
- onnx-O3：ONNX 使用 float32 精度和 O3 优化，通过 export_optimized_onnx_model(..., "O3", ...) 和 backend="onnx"。
- onnx-O4：ONNX 使用 float16 精度和 O4 优化，通过 export_optimized_onnx_model(..., "O4", ...) 和 backend="onnx"。
- onnx-qint8：ONNX 量化为 int8，使用“avx512_vnni”，通过 export_dynamic_quantized_onnx_model(..., "avx512_vnni", ...) 和 backend="onnx"。不同的量化配置产生了大致相当的加速效果。
- openvino：OpenVINO，通过 backend="openvino"。
- openvino-qint8：OpenVINO 量化为 int8，通过 export_static_quantized_openvino_model(..., OVQuantizationConfig(), ...) 和 backend="openvino"。

请注意，模型、数据集和批次大小的激进平均可能会隐藏一些更复杂的模式。例如，ONNX 在低批次大小下表现似乎更强。然而，ONNX 和 OpenVINO 甚至可能比 PyTorch 表现稍差，因此我们建议您使用特定模型和数据测试不同的后端，以找到最适合您用例的后端。

建议

根据基准测试，此流程图应能帮助您决定为模型使用哪个后端：

        %%{init: {
   "theme": "neutral",
   "flowchart": {
      "curve": "bumpY"
   }
}}%%
graph TD
A("What is your hardware?") -->|GPU| B("Are you using a small<br>batch size?")
A -->|CPU| C("Are you open to<br>quantization?")
B -->|yes| D[onnx-O4]
B -->|no| F[float16]
C -->|yes| G[openvino-qint8]
C -->|no| H("Do you have an Intel CPU?")
H -->|yes| I[openvino]
H -->|no| J[onnx]
click D "#optimizing-onnx-models"
click F "#pytorch"
click G "#quantizing-openvino-models"
click I "#openvino"
click J "#onnx"

注意

您的实际效果可能会有所不同，您应该始终使用特定模型和数据测试不同的后端，以找到最适合您用例的后端。

用户界面

这个 Hugging Face Space 提供了一个用户界面，用于导出、优化和量化 ONNX 或 OpenVINO 模型。

sentence-transformers/backend-export