最近总有学妹问我，论文要截稿了，模型来不及跑怎么办？
还有学妹问我，有好多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 LSTransformerEncoderLayer def 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, 80 inputs = 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，让更多的人享受到极速的快乐。
