GTE Chinese Embedding模型转换为Openvino IR模型

1、概述

1.1 目标

本文档描述了将 HuggingFace 格式的 nlp_gte_sentence-embedding_chinese-base 中文文本嵌入模型转换为 OpenVINO IR 格式的完整流程，以实现：

• 性能优化: 在 Intel CPU/GPU 上获得更快的推理速度

• 体积优化: 减小模型部署包大小（从 587 MB 降至 196 MB）

• 跨平台部署: 便于在各种设备上部署

1.2 转换路径

PyTorch 模型 (HuggingFace)
    ↓
ONNX 模型 (中间格式)
    ↓
OpenVINO IR (最终格式)

1.3 关键指标对比

1.4 Python 依赖

PyTorch → ONNX 阶段:

transformers>=4.30.0
torch>=2.0.0
sentencepiece>=0.1.99
onnx>=1.14.0
onnxruntime>=1.15.0
numpy>=1.24.0

ONNX → OpenVINO 阶段:

openvino>=2024.0.0
transformers>=4.30.0
numpy>=1.24.0
sentencepiece>=0.1.99

2、环境准备

2.1 创建 Python 虚拟环境

# 创建虚拟环境
conda create -n ov_convert python=3.11

# 激活虚拟环境
conda activate ov_convert

2.2 安装依赖

pip install transformers torch onnx onnxruntime openvino numpy sentencepiece
# 可选
pip install onnxoptimizer

3、转换流程

3.1 阶段一：PyTorch → ONNX

3.1.1 转换原理

ONNX (Open Neural Network Exchange) 是一个开放的深度学习模型表示格式，支持在不同框架间转换模型。

转换步骤:

1. 加载 PyTorch 模型

2. 生成示例输入（用于推导图结构）

3. 使用 torch.onnx.export() 导出

4. 应用 ONNX 优化器

5. 验证模型正确性

3.1.2 转换 ONNX 模型

from pathlib import Path
import numpy as np

def export_to_onnx(
    model_dir: str,
    output_path: str,
    opset_version: int = 18,
    optimize: bool = True,
    verify: bool = True,
):
    import torch
    from transformers import AutoModel, AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
    base_model = AutoModel.from_pretrained(model_dir, trust_remote_code=True)
    base_model = base_model.to(torch.float32)
    base_model.eval()

    class LastHiddenStateModel(torch.nn.Module):
        def __init__(self, m):
            super().__init__()
            self.m = m

        def forward(self, input_ids, attention_mask):
            return self.m(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state

    model = LastHiddenStateModel(base_model).eval()

    if opset_version < 18:
        opset_version = 18

    sample = tokenizer(
        ["这是一个示例文本用于转换模型"],
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=512,
    )

    out = Path(output_path)
    out.parent.mkdir(parents=True, exist_ok=True)

    torch.onnx.export(
        model,
        (sample["input_ids"], sample["attention_mask"]),
        str(out),
        input_names=["input_ids", "attention_mask"],
        output_names=["last_hidden_state"],
        dynamic_axes={
            "input_ids":