只用几行代码，我让模型『训练』加速了3倍以上！

❤️点击上方，选择星标或置顶，每天给你送上干货❤️

最近总有学妹问我，论文要截稿了，模型来不及跑怎么办？

还有学妹问我，有好多idea，验证一个就要跑一周怎么办？

这时候我想起了下面这张图，我想这句话反映了大多数从事人工智能的科研工作者们目前的状态。

于是我告诉学妹，想要快，找我就对了，我教你怎样让你的模型训练加速3倍以上！

这里我们需要用到的就是字节跳动AI Lab最近开源的「新版训练加速引擎」——LightSeq。

具体的原理这里就不做过多介绍了，过两天会专门发布一篇介绍技术细节的文章，敬请期待。今天我来教大家如何使用LightSeq。

安装步骤

源码安装

你可以从源码进行安装，使用如下命令：

git clone https://github.com/bytedance/lightseq.git
cd lightseq
pip install -e .

如果你想执行LightSeq提供的现成样例，或者使用它的单元测试工具，那最好从源码安装。

pip安装

当然如果你想直接调用LightSeq的接口，不需要它的样例或者单元测试工具，我更推荐你用下面pip的方式安装，更加方便：

pip install lightseq

使用教程

自定义模型

大多数同学可能想自己搭建一个Transformer模型，然后用来训练各种数据，那我这里就教大家如何快速搭建一个LightSeq版本的Transformer编码层。

你只需要创建一个配置对象，然后用它创建LightSeq的编码层即可。

我写了一份完整的训练代码，非常浅显易懂，看注释就行了，亲测可以直接运行哦：

import torch
from lightseq.training.ops.pytorch.transformer_encoder_layer import LSTransformerEncoderLayerdef train(model, inputs, masks):inputs = inputs.to(device="cuda:0")masks = masks.to(device="cuda:0")model.to(device="cuda:0")model.train()opt = torch.optim.Adam(model.parameters(), lr=1e-3)for epoch in range(1000):opt.zero_grad()outputs = model(inputs, masks)loss = torch.square(outputs).mean()loss.backward()opt.step()if epoch % 200 == 0:print("epoch {:>3d}: loss = {:>5.3f}".format(epoch, loss))if __name__ == "__main__":# 定义LightSeq配置config = LSTransformerEncoderLayer.get_config(max_batch_tokens=4096,max_seq_len=256,hidden_size=1024,intermediate_size=4096,nhead=16,attn_prob_dropout_ratio=0.1,activation_dropout_ratio=0.1,hidden_dropout_ratio=0.1,pre_layer_norm=True,fp16=False,local_rank=0)# 随机生成输入bsz, sl = 10, 80inputs = torch.randn(bsz, sl, config.hidden_size)masks = torch.zeros(bsz, sl)# 定义LightSeq编码层并进行训练model = LSTransformerEncoderLayer(config)train(model, inputs, masks)

下面两个样例都放在了lightseq/training/examples目录下，推荐大家采用源码安装的方式安装LightSeq，这样可以直接运行样例。

Hugging Face

Hugging Face是目前用的最多的预训练模型库了吧，主要是用起来太方便了，模型也很全。直接pip install transformers安装即可。

以BERT在NER任务上微调为例，直接运行LightSeq提供的脚本就行：

sh lightseq/training/examples/huggingface/run_ner.sh

Fairseq

Fairseq是目前最主流的序列生成库之一，用来做机器翻译、文本生成等任务都是非常方便的。安装的话也很简单，直接pip install fairseq即可。

LightSeq同样提供了现成的运行脚本，如果想运行LightSeq加速后的模型，执行如下命令：

sh lightseq/training/examples/fairseq/ls_fairseq_wmt14en2de.sh

再来看看细节，一般如果我们想用Fairseq来训练一个机器翻译模型，通常首先会准备好数据集，然后执行如下命令：

fairseq-train DATA_DIR \--arch transformer_wmt_en_de_big_t2t \--optimizer adam \--criterion label_smoothed_cross_entropy \...

注意这里我们只列出了同LightSeq有关的三个参数：--arch、--optimizer和--criterion，分别指定了模型结构、参数优化器和损失函数。

如果想用LightSeq进行加速，直接将上面的运行命令改为下面这样：

lightseq-train DATA_DIR \--arch ls_transformer_wmt_en_de_big_t2t \--optimizer ls_adam \--criterion ls_label_smoothed_cross_entropy \...

注意改动的地方有4个。fairseq-train改成lightseq-train，这是为了导入LightSeq的目录。--arch、--optimizer和--criterion都加上了ls_前缀，这样就快速替换为了LightSeq的组件。

训练速度

说了这么多，实际速度到底怎么样？我用Fairseq测了一下训练的总耗时：

不同模型大小、不同批处理大小、不同显卡上加速效果都是有区别的，但总体上都能缩短一半左右的训练时间。

如果你的显卡比较老旧（我相信大多数学校实验室都是这样的），显存又比较小，那么批处理大小只能设置的很小，那加速比甚至能达到3倍以上。

项目地址

学妹试了直叫好，说用起来确实快。

你也别忘了点个star，让更多的人享受到极速的快乐。

「LightSeq地址：」
https://github.com/bytedance/lightseq

- END -

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

公众号后台回复【内推】

可以通过我的内推码投递简历，加我微信还能随时查进度、咨询问题。

公众号后台回复【加群】

可以进我的技术交流群和内推群。

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