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

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

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

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

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

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

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

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

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

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

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

https://huggingface.co/transformers/model_doc/bart.html

huggingface bart

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

from transformers import BartTokenizer, BartForConditionalGeneration
tokenizer = BartTokenizer.from_pretrained("facebook/bart-base")
model = BartForConditionalGeneration.from_pretrained("facebook/bart-base")
sentences = ["I love that girl, but <mask> does not <mask> 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

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

https://github.com/godweiyang/lightseq/tree/python_example/example/python

运行前要先安装一下lightseq包:

pip3 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 lightseq
from transformers import BartTokenizer
tokenizer = BartTokenizer.from_pretrained("facebook/bart-base")
model = lightseq.Transformer("lightseq_bart_base.pb", 128)
sentences = ["I love that girl, but <mask> does not <mask> 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 <mask> does not <mask> me.
She is so <mask> that I can not help glance at <mask>.
Nothing's gonna <mask> my love for you.
Drop everything now. Meet me in the pouring <mask>. Kiss me on the sidewalk.

运行结果如下:

=========================lightseq=========================
lightseq generating...
lightseq time: 0.034502994269132614s
lightseq 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.6297104470431805s
huggingface 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.Transformermodel.infer进行快速推理。

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

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

我又甩给她一串地址:

https://github.com/bytedance/lightseq

速度超快!字节跳动开源序列推理引擎LightSeq

好好看,好好学,都是CUDA写的,要是看得迷糊,建议先去看看我之前的入门教程嗷:

熬了几个通宵,我写了份CUDA新手入门代码

从此,世上又多了一位快如⚡️的👧。

相关文章
|
3月前
英伟达玩转剪枝、蒸馏:把Llama 3.1 8B参数减半,性能同尺寸更强
【9月更文挑战第10天】《通过剪枝和知识蒸馏实现紧凑型语言模型》由英伟达研究人员撰写,介绍了一种创新方法,通过剪枝和知识蒸馏技术将大型语言模型参数数量减半,同时保持甚至提升性能。该方法首先利用剪枝技术去除冗余参数,再通过知识蒸馏从更大模型转移知识以优化性能。实验结果显示,该方法能显著减少模型参数并提升性能,但可能需大量计算资源且效果因模型和任务而异。
83 8
|
3月前
|
机器学习/深度学习 存储 人工智能
压缩大型语言模型(LLMs):缩小10倍、性能保持不变
尽管大规模语言模型(LLMs)在多种应用场景中表现出色,但其庞大的规模也带来了实际部署难题。本文探讨了通过模型压缩技术解决这些问题的方法,介绍了量化、剪枝和知识蒸馏三种主要压缩技术,并通过具体Python代码示例展示了如何将一个100M参数的文本分类模型压缩至52.8M参数,再通过4位量化进一步减小至原来的1/7,同时保持甚至提升性能。示例代码展示了从数据预处理、模型训练到评估的完整流程,证明了压缩技术的有效性。
152 6
|
2月前
|
人工智能 网络架构 开发者
第一个100%开源的MoE大模型,7B的参数,1B的推理成本
【10月更文挑战第1天】《OLMoE: Open Mixture-of-Experts Language Models》一文介绍了OLMoE,这是一个完全开源的Mixture-of-Experts(MoE)语言模型,具有70亿参数量,但每个输入令牌仅需10亿参数进行推理,有效平衡了性能与成本。OLMoE由Allen Institute for AI等机构的研究者共同开发,其开源特性促进了大规模语言模型领域的合作与创新,降低了资源浪费,同时具备良好的可扩展性和效率,为研究和应用开辟了新可能。然而,其复杂性也可能带来训练和调优上的挑战。
52 2
|
7月前
|
机器学习/深度学习 自然语言处理 测试技术
SUPRA:无须额外训练,将Transformer变为高效RNN,推理速度倍增
`Transformers`模型因其在多种任务上的优秀性能而广泛采用,但其内存需求和推理成本随词元数量指数增长。为解决这一问题,论文《Linearizing Large Language Models》提出`SUPRA`方法,将预训练的`Transformers`转换为递归神经网络(RNN),实现有效推理并降低训练成本至原来的5%。`SUPRA`通过将注意力机制线性化并用`GroupNorm`替代`softmax`,保持预训练模型的优势。经过微调,转换后的模型在标准语言理解和长上下文任务中保持高性能,展示了在长序列处理方面的潜力,但仍有改进空间。
156 2
|
7月前
|
机器学习/深度学习 并行计算 PyTorch
【多GPU炼丹-绝对有用】PyTorch多GPU并行训练:深度解析与实战代码指南
本文介绍了PyTorch中利用多GPU进行深度学习的三种策略:数据并行、模型并行和两者结合。通过`DataParallel`实现数据拆分、模型不拆分,将数据批次在不同GPU上处理;数据不拆分、模型拆分则将模型组件分配到不同GPU,适用于复杂模型;数据和模型都拆分,适合大型模型,使用`DistributedDataParallel`结合`torch.distributed`进行分布式训练。代码示例展示了如何在实践中应用这些策略。
2003 2
【多GPU炼丹-绝对有用】PyTorch多GPU并行训练:深度解析与实战代码指南
|
人工智能 自然语言处理 测试技术
只用几行代码,我让模型『训练』加速了3倍以上!
只用几行代码,我让模型『训练』加速了3倍以上!
161 0
只用几行代码,我让模型『训练』加速了3倍以上!
|
机器学习/深度学习 存储 人工智能
神经网络推理加速入门——转换提前
一个很简单的转换提前,就能让整个神经网络的性能提升一倍,甚至跟高!
神经网络推理加速入门——转换提前
|
机器学习/深度学习 传感器 编解码
全新剪枝框架 | YOLOv5模型缩减4倍,推理速度提升2倍(一)
全新剪枝框架 | YOLOv5模型缩减4倍,推理速度提升2倍(一)
470 0
|
存储 编解码 算法
全新剪枝框架 | YOLOv5模型缩减4倍,推理速度提升2倍(二)
全新剪枝框架 | YOLOv5模型缩减4倍,推理速度提升2倍(二)
490 0
|
机器学习/深度学习 自然语言处理
微软提出MiniViT | 把DeiT压缩9倍,性能依旧超越ResNet等卷积网络(二)
微软提出MiniViT | 把DeiT压缩9倍,性能依旧超越ResNet等卷积网络(二)
253 0