• Stars
    star
    145
  • Rank 254,144 (Top 6 %)
  • Language
    Python
  • License
    MIT License
  • Created over 2 years ago
  • Updated over 2 years ago

Reviews

There are no reviews yet. Be the first to send feedback to the community and the maintainers!

Repository Details

sentence-transformers to onnx 让sbert模型推理效率更快

quick_sentence_transformers

convert sentence-transformer to onnx model

sentence-transformer推理加速

背景

因为工作需要,做到需要使用nlp问答的内容,然后就考虑到sentence-transformer(以下简称为sbert模型)。而且sbert句子转向量这个方法感觉很高效,因此就考虑到这个维度。

但是在使用sbert的过程中,遇到推理速度问题,也就是sbert的encode速度太慢了。后来花费了大量的时间在提高模型推理效率的工作上。

最后摸索出一个推理的高效方法:将模型转换成onnx格式的文件或者转换成tensorrt的plan格式文件。

面临的问题:

这个模型不能直接转换么?确实不行。

在作者的一个issue可以找到答案:

UKPLab/sentence-transformers#46

Sadly I never worked with onnx. In SentenceTransformer, the forward function takes in one argument: features (and the second in python is self). Features is a dictionary, that contains the different features, for example, token ids, word weights, attention values, token_type_ids. For the BERT model, I think your input must look like this:

input_features = {'input_ids': dummy_input0, 'token_type_ids': dummy_input1, 'input_mask': dummy_input2}

And then:

torch.onnx.export(model,input_features, onnx_file_name, verbose=True)

在这个链接下,作者说他的模型已经放在了huggingface Models Hub上,且Huggingface支持转换onnx。我尝试了一下,还是会出现问题。

解决办法

既然整个的sbert模型不能直接转换,而且sbert中非常耗时间的模型其实是在transformer部分。那么我单独对transformer部分做加速不行么?答案是可以的。

为什么会想到要将sbert分开加速

以sbert的paraphrase-multilingual-MiniLM-L12-v2模型为案例。模型可以从这个链接下载https://public.ukp.informatik.tu-darmstadt.de/reimers/sentence-transformers/v0.2/

这个模型解压之后,模型实际上分为两个部分:

  1. 第一部分是0_Transformer部分
  2. 第二部分是1_Pooling部分。

可是两个模型的导入,并没有在sbert包里面明显的体现出来。使用的是import_from_string方式,非常巧妙。说白了:就是基于importlib包和模型文件里面的modules.json配合,巧妙的导入模型: 具体代码如下:

import importlib
import os
import json
from collections import OrderedDict


def import_from_string(dotted_path):
    """
    Import a dotted module path and return the attribute/class designated by the
    last name in the path. Raise ImportError if the import failed.
    """
    try:
        module_path, class_name = dotted_path.rsplit('.', 1)
    except ValueError:
        msg = "%s doesn't look like a module path" % dotted_path
        raise ImportError(msg)

    try:
        module = importlib.import_module(dotted_path)
    except:
        module = importlib.import_module(module_path)

    try:
        return getattr(module, class_name)
    except AttributeError:
        msg = 'Module "%s" does not define a "%s" attribute/class' % (module_path, class_name)
        raise ImportError(msg)
        
        
        
model_path = "../models/paraphrase-multilingual-MiniLM-L12-v2"



modules_json_path = os.path.join(model_path, 'modules.json')
with open(modules_json_path) as fIn:
    modules_config = json.load(fIn)

modules = OrderedDict()
for module_config in modules_config:
    module_class = import_from_string(module_config['type'])
    module = module_class.load(os.path.join(model_path, module_config['path']))
    modules[module_config['name']] = module

短短几行代码,就可以把两个模型都导入。而这两个模型,都是相互独立的,在推理的时候,就是把transforemr的token_embeddings和tokenizer后的attention_mask输入到pooling模型里面。

具体操作

具体的操作步骤其实非常简单了,就是把第一部分的模型转换一下,然后再把模型的输出和pooling结合起来就可以。

代码太多了,我这里就不分享了。直接给大家仓库地址: https://github.com/yuanzhoulvpi2017/quick_sentence_transformers

  1. 具体的转换步骤,可以查看 仓库里面的03_sbert2onnx_gpu.ipynb
  2. 转换后的模型,如何用在pooling层上呢,可以查看仓库里面的06speedtest.py

感谢

中间参考了大量资料,列举如下:

  1. https://towardsdatascience.com/hugging-face-transformer-inference-under-1-millisecond-latency-e1be0057a51c 这个链接给我很大的影响,让我知道tensorrt和onnx推理能比pytorch直接跑快那么多,但是现在triton我还没搞,估计这个肯定比fastapi要快不少。

  2. https://github.com/shibing624/text2vec 感谢徐明大佬的text2vec。我就是在阅读他的仓库的时候,发现了sentence-transformer包。同时也受到他这个包的一些写法的启发。

  3. https://maple.link/2021/06/10/ONNX%E6%A8%A1%E5%9E%8B%E8%BD%ACTRT%20engine/ 我是从这个大佬的文章里面看到,如何把模型从onnx转换成tenosrrt的。写的非常清楚。后来改了他的代码一些细节部分。

  4. https://onnxruntime.ai/docs/tutorials/huggingface.html 这个是onnxruntime官网,我在这个官网案例了,找到了如何加速transofmer,进而学会了如何加速sbert的transformer部分。

  5. https://github.com/UKPLab/sentence-transformers 说到底还是看sbert作者本身的代码,我从这个包里面学到了很多知识,不限于python、pytorch、数据流的控制、loss修改等部分,学习到非常非常多新的思想。

More Repositories

1

zero_nlp

中文nlp解决方案(大模型、数据、模型、训练、推理)
Jupyter Notebook
2,573
star
2

tiny_python

tiny python project
Jupyter Notebook
154
star
3

DocumentSearch

基于sentence transformers和chatglm实现的文档搜索工具
Python
150
star
4

vscode_debug_transformers

Python
106
star
5

WechatGroupRobot

微信群聊机器人,微信,chatgpt,python,自动化
Python
91
star
6

SentenceEmbedding

Jupyter Notebook
61
star
7

questionAnswerSystem

基于sentence-transformers实现文本转向量的机器人
Jupyter Notebook
45
star
8

TimeSeries

Jupyter Notebook
27
star
9

plot_data

这个是我画图的一些数据和代码还有结果
Jupyter Notebook
25
star
10

chinamap

Python
25
star
11

mamba4transformers

Python
12
star
12

nlpboss

Python
6
star
13

Rust4SenVec

convert sentence to vector by nlp transformers model in Rust
Jupyter Notebook
6
star
14

yuanzhoulvpi2017

personal info
5
star
15

countdown

Rust
4
star
16

newvennplot

Jupyter Notebook
4
star
17

r_practice

这个是我给一些人写的代码,保存一下,还有数据
R
3
star
18

getweather

我的第一个包,用来查询中国城市(包括县)的天气
R
3
star
19

getdata

Jupyter Notebook
3
star
20

master_shiny_CN

2
star
21

mytools

R
1
star
22

sklearn

自己从sklearn网上找到的,种类很多,如果是入门者,可以每天检查写这个代码两个以上,坚持写完,必有成效
Python
1
star
23

tiny_pytorch

从网上找到pytorch一些代码,然后自己抄写,并且修正那些代码,并且发布到这上面来
Python
1
star
24

R4pi

C++
1
star
25

raincloud_plot

use R create raincloud
HTML
1
star
26

tiny_cpp

C++
1
star
27

MakeTextLabel

Python
1
star
28

classdigits_inter

Interactive handwritten digit recognition
Python
1
star
29

tiny_tensorflow

准备筹集tensorflow代码,从入门到进阶
Python
1
star
30

dendrogram_maps

Jupyter Notebook
1
star
31

gongtongqujian

就是计算区间重合的最大数目
Python
1
star
32

down_issue

R
1
star