只用两行代码,我让Transformer推理加速了50倍

网友投稿 1064 2022-11-13

只用两行代码,我让Transformer推理加速了50倍

只用两行代码,我让Transformer推理加速了50倍

最近有学妹问我,我训了一个Transformer模型,但是预测好慢啊,有啥解决方案吗?

我心想,你又想好,又想快,咋不上天-?呢?

于是我跟她说,你可以试试lightseq啊,跟闪电⚡️一样快,用了你就可以上天了。

她一脸懵比,lightseq是啥玩意儿啊?咋就能让我的模型起飞-?️了呢?

我跟她说,你不需要知道太多细节,你只需要知道它是一个Transformer系列模型推理加速库就行了。

她还是一脸疑惑,那用起来能有huggingface方便吗?你看人家就两行代码。

我不屑一笑,就这?lightseq也只要两行代码就够了!

为了方便,我用了一个bart模型预测句子中mask单词的例子来给她吹了一波。

不懂什么是bart?建议先去看看huggingface的文档:​​​bart

我们平时想用huggingface的bart来预测句子中的mask单词,大体上都会像下面这样写代码:

from transformers import BartTokenizer, BartForConditionalGenerationtokenizer = BartTokenizer.from_pretrained("facebook/bart-base")model = BartForConditionalGeneration.from_pretrained("facebook/bart-base")sentences = ["I love that girl, but does not me."]inputs = tokenizer(sentences, return_tensors="pt", padding=True)generated_ids = model.generate(inputs["input_ids"], max_length=50)res = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)print(res)

当然运行前要先安装一下​​transformers​​包:

pip3 install transformers

最后会输出句子“I love that girl, but she does not love me.”,句子中的两个“mask”被预测成了“she”和“love”。

看起来预测的很nice,但是预测的也太慢了,这要是有一堆句子要去预测,不得等到-?年-?月?

接下来我们来看看lightseq是怎么加速预测的。

lightseq bart

代码我都放在下面地址了,只要两分钟就能跑出结果了:​​​install lightseq

首先lightseq只能接收Protocol Buffer协议定义的模型文件,如果你不知道这是啥也没关系,因为我们帮你写好了模型转换的脚本,就是​​hf_bart_export.py​​​,它会将huggingface预训练的bart模型转换为​​transformer_pb2.py​​定义好的Protocol Buffer格式。

所以直接运行​​python3 hf_bart_export.py​​就行了,这里我们用的是bart-base模型。

运行完了会发现执行目录下多出一个​​lightseq_bart_base.pb​​文件,这就是转换后的模型文件。

最后直接跟huggingface一样,两行代码就能搞定啦:

import lightseqfrom transformers import BartTokenizertokenizer = BartTokenizer.from_pretrained("facebook/bart-base")model = lightseq.Transformer("lightseq_bart_base.pb", 128)sentences = ["I love that girl, but does not me."]inputs = tokenizer(sentences, return_tensors="pt", padding=True)generated_ids = model.infer(inputs["input_ids"])generated_ids = [ids[0] for ids in generated_ids[0]]res = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)print(res)

看得出来仅仅替换了模型定义和模型推理那两行代码而已,是不是非常简单快速?

这时候她又问了,那我换一个模型,比如bert,要怎么导出pb模型呢?

也很简单,只需要为bert也单独写一个​​hf_bert_export.py​​就行了。不过目前还在开发中,之后会慢慢完善常见的一些模型的。

速度到底怎么样?

我写好了一个例子,就在​​ls_bart.py​​​里,直接运行就行了,当然你也可以加上​​--user_input​​参数来手动输入句子。

输入的句子是:

I love that girl, but does not me.She is so that I can not help glance at .Nothing's gonna my love for you.Drop everything now. Meet me in the pouring . Kiss me on the sidewalk.

运行结果如下:

=========================lightseq=========================lightseq generating...lightseq time: 0.034502994269132614slightseq results:I love that girl, but she does not love me.She is so beautiful that I can not help glance at her.Nothing's gonna change my love for you.Drop everything now. Meet me in the pouring rain. Kiss me on the sidewalk.=========================huggingface=========================huggingface generating...huggingface time: 1.6297104470431805shuggingface results:I love that girl, but she does not love me.She is so beautiful that I can not help glance at her.Nothing's gonna change my love for you.Drop everything now. Meet me in the pouring rain. Kiss me on the sidewalk.

可以看出预测的是真的--??,最后两句歌词都预测的很完美,能看出是啥歌吗?

再看预测时间,lightseq是huggingface的47倍左右,真是一个天上一个地下啊。

总结

总结一下,想要使用lightseq加速你的模型,只需要两步就行了:

将你的模型转换为pb格式的模型。(lightseq为你写好了转换脚本,不断更新中)调用​​lightseq.Transformer​​和​​model.infer​​进行快速推理。

学妹赶紧打住了我,好了好了,我知道很--??了。还给你装起来了,我这就去用。

但是源码哪里有?我想学一学。

我又甩给她一串地址:​​​END -

我是godweiyang,华东师范大学计算机系本硕专业第一,字节跳动AI Lab NLP算法工程师,秋招斩获上海三家互联网大厂ssp offer,主要研究方向为机器翻译、句法分析、模型压缩与加速。最大特点就是脾气好、有耐心,有任何问题都可以随时咨询我,不管是技术上的还是生活上的。

可以进我的技术交流群

记得一键③连,今天的你格外的可爱-?

版权声明:本文内容由网络用户投稿,版权归原作者所有,本站不拥有其著作权,亦不承担相应法律责任。如果您发现本站中有涉嫌抄袭或描述失实的内容,请联系我们jiasou666@gmail.com 处理,核实后本网站将在24小时内删除侵权内容。

上一篇:MongoDB第二篇
下一篇:【NLP 算法岗】提前批暑期实习面(试)经(历)
相关文章

 发表评论

暂时没有评论,来抢沙发吧~