BGE推理加速实战：从零搭建高性能文本嵌入服务

1次阅读

共计 1790 个字符，预计需要花费 5 分钟才能阅读完成。

BGE（Bidirectional Generative Encoder）模型在文本检索、语义聚类等场景中表现出色，能有效捕捉上下文语义信息。但原生 PyTorch 实现存在推理延迟高、显存占用大的问题，尤其处理长文本时性能下降明显。本文将带您从零实现一套完整的推理加速方案，涵盖模型转换、量化优化和工程化部署全流程。

原始 PyTorch 推理
优点：开发简单，直接加载.h5 或.pt 模型文件
缺点：计算图解释执行开销大，无法利用硬件加速
ONNX Runtime 优化
支持算子融合 (Operator Fusion) 和常量折叠
自动选择最优执行提供者（CUDA/DNNL）
典型加速比：1.3-1.8x
TensorRT 部署
通过层融合和内核自动调优实现极致优化
需要处理动态 shape 的额外配置
典型加速比：2-3x
量化方案
FP16：保持 90%+ 精度，显存减半
INT8：需校准数据集，精度损失约 2 -5%
典型加速比：FP16 1.5x, INT8 2-3x

# 导出时指定动态维度（batch_size 和 seq_len）dynamic_axes = {'input_ids': {0: 'batch_size', 1: 'sequence_length'},
    'attention_mask': {0: 'batch_size', 1: 'sequence_length'}
}

torch.onnx.export(
    model,
    dummy_input,
    "bge.onnx",
    input_names=["input_ids", "attention_mask"],
    output_names=["embeddings"],
    dynamic_axes=dynamic_axes,  # 关键配置！opset_version=13
)

class BGEService:
    def __init__(self, onnx_path):
        self.ort_session = ort.InferenceSession(onnx_path)

    def batch_infer(self, texts: List[str]):
        # 统一 padding 到本 batch 最大长度
        inputs = tokenizer(texts, padding=True, return_tensors="np")

        # GPU 内存监控（需安装 pynvml）handle = nvml.nvmlDeviceGetHandleByIndex(0)
        info = nvml.nvmlDeviceGetMemoryInfo(handle)
        print(f"GPU 内存占用：{info.used/1024**2:.2f}MB")

        # 执行推理
        outputs = self.ort_session.run(
            None,
            {"input_ids": inputs["input_ids"], 
             "attention_mask": inputs["attention_mask"]}
        )
        return outputs[0]  # [batch_size, hidden_dim]