加速推理

Sentence Transformers 支持 3 种后端来计算使用稀疏编码器模型的稀疏嵌入，每种后端都有其自身的优化来加速推理。

稀疏编码器的默认后端。灵活高效的模型加速器。模型优化，主要用于英特尔硬件。不同后端的基准测试。用于导出、优化和量化模型的图形用户界面。

PyTorch

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

from sentence_transformers import SparseEncoder

model = SparseEncoder("naver/splade-v3")

sentences = ["This is an example sentence", "Each sentence is converted"]
embeddings = model.encode(sentences)

decoded = model.decode(embeddings)
print(decoded[0][:5])
# [('example', 2.451861619949341), ('sentence', 2.214038848876953), ('examples', 2.0835916996002197), ('sentences', 2.0063159465789795), ('this', 1.7662484645843506)]

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

float16 (fp16)

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

from sentence_transformers import SparseEncoder

model = SparseEncoder("naver/splade-v3", model_kwargs={"torch_dtype": "float16"})
# or: model.half()

sentences = ["This is an example sentence", "Each sentence is converted"]
embeddings = model.encode(sentences)

bfloat16 (bf16)

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

from sentence_transformers import SparseEncoder

model = SparseEncoder("naver/splade-v3", model_kwargs={"torch_dtype": "bfloat16"})
# or: model.bfloat16()

sentences = ["This is an example sentence", "Each sentence is converted"]
embeddings = model.encode(sentences)

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 SparseEncoder

model = SparseEncoder("naver/splade-v3", backend="onnx")

sentences = ["This is an example sentence", "Each sentence is converted"]
embeddings = model.encode(sentences)

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

注意

如果您希望在 Sentence Transformers 之外使用 ONNX 模型，您需要自己执行 SPLADE 池化和/或归一化。ONNX 导出仅转换（掩码语言建模）Transformer 组件，该组件输出词元嵌入，而不是全文的稀疏嵌入。

所有通过 model_kwargs 传递的关键字参数都将传递给 ORTModelForMaskedLM.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 = SparseEncoder("path/to/my/model", backend="onnx")
model.save_pretrained("path/to/my/model")

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

model = SparseEncoder("naver/splade-v3", backend="onnx")
model.push_to_hub("naver/splade-v3", create_pr=True)

优化 ONNX 模型

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

model: 一个使用 ONNX 后端加载的 Sentence Transformer、Sparse Encoder 或 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 时创建一个拉取请求（pull request）。当您没有仓库的写权限时很有用。
file_suffix：（可选）在保存模型时附加到模型名称的字符串。如果未指定，将使用优化级别的名称字符串，或者如果优化配置不只是一个字符串优化级别，则仅使用 "optimized"。

请参阅此示例，了解如何使用优化级别 3（基本和扩展通用优化、transformers 特定融合、快速 Gelu 近似）导出模型：

Hugging Face Hub 模型

仅优化一次

from sentence_transformers import SparseEncoder, export_optimized_onnx_model

model = SparseEncoder("naver/splade-v3", backend="onnx")
export_optimized_onnx_model(
    model=model,
    optimization_config="O3",
    model_name_or_path="naver/splade-v3",
    push_to_hub=True,
    create_pr=True,
)

在拉取请求合并之前

from sentence_transformers import SparseEncoder

pull_request_nr = 2 # TODO: Update this to the number of your pull request
model = SparseEncoder(
    "naver/splade-v3",
    backend="onnx",
    model_kwargs={"file_name": "onnx/model_O3.onnx"},
    revision=f"refs/pr/{pull_request_nr}"
)

拉取请求合并后

from sentence_transformers import SparseEncoder

model = SparseEncoder(
    "naver/splade-v3",
    backend="onnx",
    model_kwargs={"file_name": "onnx/model_O3.onnx"},
)

本地模型

仅优化一次

from sentence_transformers import SparseEncoder, export_optimized_onnx_model

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

优化后

from sentence_transformers import SparseEncoder

model = SparseEncoder(
    "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、Sparse Encoder 或 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 时创建一个拉取请求（pull request）。当您没有仓库的写权限时很有用。
file_suffix：（可选）在保存模型时附加到模型名称的字符串。如果未指定，将使用 "qint8_quantized"。

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

请参阅此示例，了解如何使用 avx512_vnni 将模型量化为 int8：

Hugging Face Hub 模型

仅量化一次

from sentence_transformers import SparseEncoder, export_dynamic_quantized_onnx_model

model = SparseEncoder("naver/splade-v3", backend="onnx")
export_dynamic_quantized_onnx_model(
    model=model,
    quantization_config="avx512_vnni",
    model_name_or_path="sentence-transformers/naver/splade-v3",
    push_to_hub=True,
    create_pr=True,
)

在拉取请求合并之前

from sentence_transformers import SparseEncoder

pull_request_nr = 2 # TODO: Update this to the number of your pull request
model = SparseEncoder(
    "naver/splade-v3",
    backend="onnx",
    model_kwargs={"file_name": "onnx/model_qint8_avx512_vnni.onnx"},
    revision=f"refs/pr/{pull_request_nr}",
)

拉取请求合并后

from sentence_transformers import SparseEncoder

model = SparseEncoder(
    "naver/splade-v3",
    backend="onnx",
    model_kwargs={"file_name": "onnx/model_qint8_avx512_vnni.onnx"},
)

本地模型

仅量化一次

from sentence_transformers import SparseEncoder, export_dynamic_quantized_onnx_model

model = SparseEncoder("path/to/my/mpnet-legal-finetuned", backend="onnx")
export_dynamic_quantized_onnx_model(
   model=model, quantization_config="avx512_vnni", model_name_or_path="path/to/my/mpnet-legal-finetuned"
)

量化后

from sentence_transformers import SparseEncoder

model = SparseEncoder(
    "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 SparseEncoder

model = SparseEncoder("naver/splade-v3", backend="openvino")

sentences = ["This is an example sentence", "Each sentence is converted"]
embeddings = model.encode(sentences)

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

注意

如果您希望在 Sentence Transformers 之外使用 OpenVINO 模型，您需要自己执行 SPLADE 池化和/或归一化。OpenVINO 导出仅转换（掩码语言建模）Transformer 组件，该组件输出词元嵌入，而不是全文的稀疏嵌入。

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

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

提示

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

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

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

model = SparseEncoder("intfloat/multilingual-e5-small", backend="openvino")
model.push_to_hub("intfloat/multilingual-e5-small", create_pr=True)

量化 OpenVINO 模型

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

model: 一个使用 OpenVINO 后端加载的 Sentence Transformer、Sparse Encoder 或 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 时创建一个拉取请求（pull request）。当您没有仓库的写权限时很有用。
file_suffix：（可选）在保存模型时附加到模型名称的字符串。如果未指定，将使用 "qint8_quantized"。

请参阅此示例，了解如何使用静态量化将模型量化为 int8：

Hugging Face Hub 模型

仅量化一次

from sentence_transformers import SparseEncoder, export_static_quantized_openvino_model

model = SparseEncoder("naver/splade-v3", backend="openvino")
export_static_quantized_openvino_model(
    model=model,
    quantization_config=None,
    model_name_or_path="sentence-transformers/naver/splade-v3",
    push_to_hub=True,
    create_pr=True,
)

在拉取请求合并之前

from sentence_transformers import SparseEncoder

pull_request_nr = 2 # TODO: Update this to the number of your pull request
model = SparseEncoder(
    "naver/splade-v3",
    backend="openvino",
    model_kwargs={"file_name": "openvino/openvino_model_qint8_quantized.xml"},
    revision=f"refs/pr/{pull_request_nr}"
)

拉取请求合并后

from sentence_transformers import SparseEncoder

model = SparseEncoder(
    "naver/splade-v3",
    backend="openvino",
    model_kwargs={"file_name": "openvino/openvino_model_qint8_quantized.xml"},
)

本地模型

仅量化一次

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

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

量化后

from sentence_transformers import SparseEncoder

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

基准测试

下图显示了不同后端在 GPU 和 CPU 上的基准测试结果。结果是在 3 个数据集和多种批量大小下取平均值得到的。

展开基准测试详情

加速比

硬件：RTX 3090 GPU, i7-17300K CPU
数据集：GPU 测试使用 2000 个样本，CPU 测试使用 1000 个样本。
- sentence-transformers/stsb：平均 38.9 个字符 (SD=13.9)
- sentence-transformers/natural-questions：仅答案，平均 619.6 个字符 (SD=345.3)
- stanfordnlp/imdb：文本重复 4 次，平均 9589.3 个字符 (SD=633.4)
模型
- naver/splade-v3：1.1 亿参数；批量大小为 8、16、32 和 64。

我还对 ibm-granite/granite-embedding-30m-sparse 进行了基准测试，但它太小了，无法有效展示不同后端的增益，所以我将其从结果中排除了。
性能比：使用了相同的模型和硬件。我们将性能与使用 fp32 的 PyTorch（即默认后端和精度）进行比较。

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

后端
- 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(..., optimization_config="O1", ...) 和 backend="onnx" 实现。
- onnx-O2: ONNX，使用 float32 精度和 O2 优化，通过 export_optimized_onnx_model(..., optimization_config="O2", ...) 和 backend="onnx" 实现。
- onnx-O3: ONNX，使用 float32 精度和 O3 优化，通过 export_optimized_onnx_model(..., optimization_config="O3", ...) 和 backend="onnx" 实现。
- onnx-O4: ONNX，使用 float16 精度和 O4 优化，通过 export_optimized_onnx_model(..., optimization_config="O4", ...) 和 backend="onnx" 实现。
- onnx-qint8：ONNX 使用“avx512_vnni”量化到 int8，通过 export_dynamic_quantized_onnx_model(..., quantization_config="avx512_vnni", ...) 和 backend="onnx" 实现。不同的量化配置带来了大致相当的加速效果。
- openvino: OpenVINO，通过 backend="openvino" 实现。
- openvino-qint8: OpenVINO，量化到 int8，通过 export_static_quantized_openvino_model(..., quantization_config=OVQuantizationConfig(), ...) 和 backend="openvino" 实现。

请注意，对模型、数据集和批量大小进行的大规模平均，使得一些更复杂的模式无法显现。例如，对于 GPU 和 CPU，ibm-granite/granite-embedding-30m-sparse 模型从各种后端获得的收益比大型模型要小。例如，在 GPU 上使用 fp16 和 bf16 平均只能带来 1.4 倍的加速。

建议

基于基准测试，以下流程图可以帮助您决定为您的模型使用哪种后端：

        %%{init: {
   "theme": "neutral",
   "flowchart": {
      "curve": "bumpY"
   }
}}%%
graph TD
A(What is your hardware?) -->|GPU| B("Are you using a<br>batch size of <= 4?")
A -->|CPU| C("Are minor performance<br>degradations acceptable?")
B -->|yes| D[onnx-O4]
B -->|no| F[bfloat16]
C -->|yes| G[openvino-qint8]
C -->|no| H(Do you have an Intel CPU?)
H -->|yes| I[openvino]
I -->|or| J[onnx-O3]
H -->|no| K[onnx-O3]
click D "#optimizing-onnx-models"
click F "#pytorch"
click G "#quantizing-openvino-models"
click I "#openvino"
click J "#optimizing-onnx-models"
click K "#optimizing-onnx-models"

注意

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

用户界面

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

sentence-transformers/backend-export