基于bce-embedding-base_v1模型将torch模型转换为onnx模型

原创

已于 2025-02-18 11:04:52 修改 · 601 阅读

1 ·

CC 4.0 BY-SA版权

文章标签：

#embedding #深度学习 #pytorch

于 2024-11-11 10:23:00 首次发布

环境信息构建：
#确定cuda和onnxruntime的版本对应关系：https://blog.51cto.com/u_15962038/12360310

conda create -n mypython311 python=3.11.4 cudnn=8.9.2.26 cudatoolkit=11.8.0

pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu118

pip install onnxruntime-gpu==1.18.1 onnx==1.17.0 numpy==1.26.4 transformers==4.46.2 triton==2.0.0

模型转换：

import os
import torch
from transformers import AutoTokenizer, AutoModel, AutoConfig
from transformers.onnx.features import FeaturesManager
BCE_MODEL_PATH='/path/bce-embedding-base_v1/'
tokenizer = AutoTokenizer.from_pretrained(BCE_MODEL_PATH)
config = AutoConfig.from_pretrained(BCE_MODEL_PATH)
bge_model = AutoModel.from_pretrained(BCE_MODEL_PATH)
bge_model.eval()
onnx_config = FeaturesManager.get_config("bert", "default")(config)
dummy_inputs = onnx_config.generate_dummy_inputs(tokenizer, framework='pt')
base_path="/path/output_model/bce-embedding-base_v1"
output_onnx_path = os.path.join(base_path, "bce-embedding-base_v1.onnx")
config_file = os.path.join(base_path, "bce-embedding-base_v1_config.json")
dynamic_axes = {
    'input_ids': {0: 'batch_size', 1: 'sequence_length'},  # 动态 batch size 和 序列长度
    'attention_mask': {0: 'batch_size', 1: 'sequence_length'},  # 动态 batch size 和 序列长度
    'token_type_ids':