加速推理

Sentence Transformers 支持 3 个后端来执行交叉编码器模型的推理，每个后端都有其自身的优化来加速推理

交叉编码器的默认后端。灵活高效的模型加速器。模型优化，主要针对 Intel 硬件。不同后端的基准测试。用于导出、优化和量化模型的 GUI。

PyTorch

PyTorch 后端是交叉编码器的默认后端。如果您不指定设备，它将使用 “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 后端，您必须安装带有 onnx 或 onnx-gpu 额外项的 Sentence Transformers，分别用于 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 中交叉编码器相同的结果。

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

provider：用于加载模型的 ONNX Runtime provider，例如 "CPUExecutionProvider"。请参阅 https://runtime.onnx.org.cn/docs/execution-providers/ 以获取可能的 provider。如果未指定，将使用最强大的 provider（例如 "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")

或者使用 model.push_to_hub() （如果您的模型来自 Hugging Face 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 模型

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

model：使用 ONNX 后端加载的 Sentence Transformer 或交叉编码器模型。
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 （基本和扩展的通用优化、特定于 transformer 的融合、快速 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 模型

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

model：使用 ONNX 后端加载的 Sentence Transformer 或交叉编码器模型。
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 中交叉编码器相同的结果。

通过 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")

或者使用 model.push_to_hub() （如果您的模型来自 Hugging Face 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 模型

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

model：使用 OpenVINO 后端加载的 Sentence Transformer 或交叉编码器模型。
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：使用 float32 精度的 PyTorch（默认）。
- torch-fp16：使用 float16 精度的 PyTorch，通过 model_kwargs={"torch_dtype": "float16"}。
- torch-bf16：使用 bfloat16 精度的 PyTorch，通过 model_kwargs={"torch_dtype": "bfloat16"}。
- onnx：使用 float32 精度的 ONNX，通过 backend="onnx"。
- onnx-O1：使用 float32 精度和 O1 优化的 ONNX，通过 export_optimized_onnx_model(..., "O1", ...) 和 backend="onnx"。
- onnx-O2：使用 float32 精度和 O2 优化的 ONNX，通过 export_optimized_onnx_model(..., "O2", ...) 和 backend="onnx"。
- onnx-O3：使用 float32 精度和 O3 优化的 ONNX，通过 export_optimized_onnx_model(..., "O3", ...) 和 backend="onnx"。
- onnx-O4：使用 float16 精度和 O4 优化的 ONNX，通过 export_optimized_onnx_model(..., "O4", ...) 和 backend="onnx"。
- onnx-qint8：量化为 int8 的 ONNX，使用 “avx512_vnni”，通过 export_dynamic_quantized_onnx_model(..., "avx512_vnni", ...) 和 backend="onnx"。不同的量化配置产生大致相当的加速。
- openvino：OpenVINO，通过 backend="openvino"。
- openvino-qint8：量化为 int8 的 OpenVINO，通过 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