小白入门：sentence-transformer 提取embedding模型转onnx

序言

本文适合小白入门，以自己训练的句子embedding模型为例，像大家展示了如何手动将sentence-transformer的模型转为onnx。

很多时候，我也不知道这段代码啥意思，但是作为应用人员，不要在意这段代码到底干了啥，除非必要。

这里不仅展示了如何转onnx，还有你部署时候，所需要的所有的文件，都打包到一个文件夹中了。

原理讲解

哪些部分可转onnx

onnx转换的时候，tokenizer部分是无法被onnx的，只有你backone模型才能进行转onnx，不要问我为啥，因为我也不知道。
我的模型使用代码如下：

from sentence_transformers import SentenceTransformer, models

# 1. backone模型，这里用的bert-small
bert_model = models.Transformer("all-MiniLM-L6-v2") 

# 2. bert_model得到的是所有单词的向量，这些向量通过pool变成一个向量，
# 再通过normalize变成单位向量，即可进行dot，计算得到cosine相似度。
pool = models.Pooling(bert_model.get_word_embedding_dimension())
normalize = models.Normalize()

# 模型组装
mymodel = SentenceTransformer(modules=[bert_model, pool, normalize])

代码区

0. 安装依赖

pip install onnx
pip install onnxruntime
pip install onnxruntime-gpu
先CPU然后GPU,不按顺序装可能会出现问题

1. 路径配置

import os

# 你自己的模型
raw_model_dir = "../model/model11_all-MiniLM-L6-v2/"
abspath, raw_model_name = os.path.split(os.path.abspath(raw_model_dir))
# onnx后，所需要的文件，都转到了这个文件夹中
onnx_dir = os.path.join(abspath, raw_model_name+"_onnx-gpu/")

if not os.path.exists(onnx_dir):
    os.mkdir(onnx_dir)
    print("build dir:", onnx_dir)

2. 测试数据

titles = ["Treehobby Metal 2PCS Front CVD Drive Shafts RC Cars Upgrade Parts for WLtoys 144001 1/14 RC Car Truck Buggy Replacement Accessories", 
"Solar System for Kids Space Toys, 8 Planets for Kids Solar System Model with Projector, Stem Educational Toys for 5 Year Old Boys Gift", 
"Bella Haus Design Peeing Gnome - 10.3 Tall Polyresin - Naughty Garden Gnome for Lawn Ornaments, Indoor or Outdoor Decorations - Red and Green Funny Flashing Gnomes", 
"FATExpress CMX500 CMX300 Parts Motorcycle CNC Front Fork Boot Shock Absorber Tube Slider Cover Gaiters for 2017 2018 2019 Rebel CMX 300 500 17-19 (Black)", 
"All Balls Racing 56-133-1 Fork and Dust Seal Kit", 
"Shaluoman Plating 5-Spoke Wheel Rims with Hard Plastic Tires for RC 1:10 Drift Car Color Black", 
"Betonex 5pcs PLASTIK MOLDS Casting Concrete Paving Garden Paths Pavement Stone Patio#S25", 
"OwnMy 5.2 Inch Rainbow Crystal Lotus Candle Tealight Holder Candlestick, Glass Votive Candle Lamps Holder Night Light Candlestick with Gift Box for Altar Windowsill Home Decor Christmas Wedding Party", 
"cnomg Pot Creative Plants DIY Container Pot Mini Fairy Garden Flower Plants and Sweet House for Decoration, Holiday Decoration, Indoor Decoration and Gift (Silver)", 
"DUSICHIN DUS-018 Foam Cannon Lance Pressure Washer Nozzle Tip Spray Gun 3000 PSI Jet Wash", 
"Haoohu Multicolored Bucket Hat for Women Men Girls Frog Fisherman Hat Beach Sun Hat for Outdoor Travel", 
"Renzline Pool CUE Glove Billiard Player - Green/Black - for Left Hand - One Size fits All",
"Hobbywing QUICRUN WP 1080 brushed (2-3S) Electronic Speed Controller Waterproof ESC With Program Box LED BEC XT60-Plug RC Car 1:10 30112750", 
"Mk Morse CSM868NTSC Metal Devil NXT Metal Cutting Circular Saw Blade, Thin Steel, 8-Inch Diameter, 68 TPI, 5/8-Inch Arbor, multi", 
"Barbie Fashionistas Doll 109", 
"KeShi Cordless Rotary Tool, Upgraded 3.7V Li-ion Rotary Accessory Kit with 42 Pieces Swap-able Heads, 3-Speed and USB Charging Multi-Purpose Power Tool for Delicate & Light DIY Small Projects", 
"White Knight 1707SBK-20AM Black Chrome M12x1.50 Bulge Acorn Lug Nut, 20 Pack", 
"Memory Foam Bath Mat Rug,16x24 Inches,Luxury Non Slip Washable Bath Rugs for Bathroom,Soft Absorbent Floor Mats of Green Leaves for Kitchen Bedroom Indoor", 
"DEWIN Airbrush Kit, Multi-purpose Airbrush Sets with Compressor -Dual Action 0.3mm 7CC Capacity Mini Air Compressor Spray Gun for Paint Makeup Tattoo Cake Decoration, Art Tattoo Nail Design", 
"Park Tool BBT-69.2 16-Notch Bottom Bracket Tool - Fits Shimano, SRAM, Chris King, Campagnolo, etc.", 
"ElaDeco 216 Ft Artificial Vines Garland Leaf Ribbon Greenery Foliage Rattan Greek Wild Jungle Decorative Accessory Wedding Party Garden Craft Wall Decoration"]

3. 准备工作

def load_all_model(path):
    #从modules.json读取模型路径
    modules_json_path = os.path.join(path, 'modules.json')
    with open(modules_json_path) as fIn:
        modules_config = json.load(fIn)
    from_backbone_path = os.path.join(path, modules_config[0].get('path'))
    from_pooling_path = os.path.join(path, modules_config[1].get('path'))
    from_Normalize_path = os.path.join(path, modules_config[2].get('path'))
    return from_backbone_path, from_pooling_path, from_Normalize_path

from_backbone_path, from_pooling_path, from_Normalize_path = load_all_model(raw_model_dir)

from transformers import AutoConfig, AutoModel, AutoTokenizer

model = AutoModel.from_pretrained(from_backbone_path)
tokenizer = AutoTokenizer.from_pretrained(from_backbone_path)

inputs = tokenizer(titles, padding=True, truncation=True, max_length=256, return_tensors="pt")

import torch

from sentence_transformers import models

pooling = models.Pooling.load(from_pooling_path)
normalize = models.Normalize.load(from_Normalize_path)

3.1迁移保存目标文件

import shutil

_dir, pooling_end_dir = os.path.split(from_pooling_path)
shutil.copytree(from_pooling_path, os.path.join(onnx_dir, pooling_end_dir))

_dir, normalize_end_dir = os.path.split(from_Normalize_path)
shutil.copytree(from_Normalize_path, os.path.join(onnx_dir, normalize_end_dir))

'../model/model11_all-MiniLM-L6-v2_onnx-gpu-test/2_Normalize'

def copy_tokenize_filename(filename):
    full_filename = os.path.join(from_backbone_path, filename)
    return shutil.copy(full_filename, os.path.join(onnx_dir, filename))
    
print(copy_tokenize_filename("tokenizer.json"))
print(copy_tokenize_filename("tokenizer_config.json"))
print(copy_tokenize_filename("vocab.txt"))

../model/model11_all-MiniLM-L6-v2_onnx-gpu-test/tokenizer.json
../model/model11_all-MiniLM-L6-v2_onnx-gpu-test/tokenizer_config.json
../model/model11_all-MiniLM-L6-v2_onnx-gpu-test/vocab.txt

4. model转onnx-gpu

device = torch.device("cuda:0")
model.eval()
model.to(device)
inputs = inputs.to(device)
export_model_path = os.path.join(onnx_dir, "model.onnx")

with torch.no_grad():
    symbolic_names = {0: 'batch_size', 1: 'max_seq_len'}
    torch.onnx.export(model,  # model being run
                      args=tuple(inputs.values()),
                      f=export_model_path,
                      opset_version=12,  # 这个值传说12比11好,当然取决于onnx和onnxruntime
                      do_constant_folding=True, 
                      input_names=['input_ids',  
                                   'attention_mask',
                                   'token_type_ids'],
                      output_names=['start', 'end'], 
                      dynamic_axes={'input_ids': symbolic_names,  
                                    'attention_mask': symbolic_names,
                                    'token_type_ids': symbolic_names,
                                    'start': symbolic_names,
                                    'end': symbolic_names})
    print("Model exported at ", export_model_path)

Model exported at  ../model/model11_all-MiniLM-L6-v2_onnx-gpu-test/model.onnx

5. 测试一下是否出错以及速度

5.1 测试速度是否OK

import onnxruntime
from torch import Tensor

export_model_path = os.path.join(onnx_dir, "model.onnx")

device = torch.device("cuda:0")

sess_options = onnxruntime.SessionOptions()
sess_options.optimized_model_filepath = export_model_path
session = onnxruntime.InferenceSession(export_model_path, sess_options, providers=['CUDAExecutionProvider']) # 你的是安装在cuda

2023-07-21 17:54:55.912264962 [W:onnxruntime:, session_state.cc:1136 VerifyEachNodeIsAssignedToAnEp] Some nodes were not assigned to the preferred execution providers which may or may not have an negative impact on performance. e.g. ORT explicitly assigns shape related ops to CPU to improve perf.
2023-07-21 17:54:55.912385419 [W:onnxruntime:, session_state.cc:1138 VerifyEachNodeIsAssignedToAnEp] Rerunning with verbose output on a non-minimal build will show node assignments.
2023-07-21 17:54:56.222846005 [W:onnxruntime:, inference_session.cc:1491 Initialize] Serializing optimized model with Graph Optimization level greater than ORT_ENABLE_EXTENDED and the NchwcTransformer enabled. The generated model may contain hardware specific optimizations, and should only be used in the same environment the model was optimized in.

pooling_gpu = pooling.cuda()
normalize_gpu = normalize.cuda()

import time
begin = time.time()
for i in range(1000):
    inputs = tokenizer(titles, padding=True, truncation=True, max_length=256, return_tensors="pt")
    ort_inputs = {k: v.cpu().numpy() for k, v in inputs.items()}
    ort_outputs = session.run(None, ort_inputs)
    ort_outputs1 = pooling_gpu.forward(features={'token_embeddings': Tensor(ort_outputs[0]),
                                           'attention_mask': Tensor(ort_inputs.get('attention_mask'))})
    ort_outputs2 = normalize_gpu.forward(ort_outputs1)['sentence_embedding']
end = time.time()    
print("cost time:", end-begin)

cost time: 31.3445

begin = time.time()
for i in range(1000):
    inputs = tokenizer(titles, padding=True, truncation=True, max_length=256, return_tensors="np")
    ort_inputs = dict(inputs)
    ort_outputs = session.run(None, ort_inputs)
    ort_outputs1 = pooling_gpu.forward(features={'token_embeddings': Tensor(ort_outputs[0]).to(device),
                                           'attention_mask': Tensor(ort_inputs.get('attention_mask')).to(device)})
    ort_outputs2 = normalize_gpu.forward(ort_outputs1)['sentence_embedding']
end = time.time()    
print("cost time:", end-begin)

cost time: 19.234

5.2测试结果是否OK

from sentence_transformers import SentenceTransformer

st_model = SentenceTransformer(raw_model_dir)
x = st_model.encode(titles)

import numpy as np
np.abs((x - ort_outputs2.cpu().numpy())).sum()

0.00010381325

误差数值很小，结果OK

6. tar 这些文件

abs_onnx_dir = os.path.abspath(onnx_dir)
# _dir, onnx_name = os.path.split(abs_onnx_dir)
os.system(f"tar -cf {abs_onnx_dir[:-1]}.tar {abs_onnx_dir}")
# f"tar -cf {abs_onnx_dir[:-1]}.tar {abs_onnx_dir}"

tar: Removing leading `/' from member names
0