roberta-tiny service

基于roberta-tiny的句向量服务

背景介绍

google发布的大规模预训练模型bert在nlp诸多任务fine-tune后刷新了SOTA，后续有许多工作跟随着bert进行开展。其中roberta以更大规模的训练语料，更久的训练时间，并在预训练阶段取消了NSP(Next Sentence Predict)，进一步提高了模型的准确度。但是无论是bert还是roberta，它们的参数量巨大（110M），对服务器gpu的资源占用率很高，模型推理速度欠佳。因此我们考虑采用更轻量级的roberta-tiny模型，开源自https://github.com/ZhuiyiTechnology/pretrained-models.

roberta-tiny

【配置】 4层模型，hidden size为384，对Embedding层做了低秩分解(384->128->384)，可以用bert4keras加载使用。

【训练】 使用bert4keras在TPU v3-8上训练，使用带梯度累积的LAMB优化器，批大小为800，累积4步更新，相当于以批大小3200训练了125k步（前3125步为warmup）。

【备注】 速度跟albert tiny一致，普通分类性能也基本一致，但由于roberta模型并没有参数共享这个约束，所以在生成式任务等复杂任务上效果优于albert tiny。

网文领域的预训练模型

为了达到更好的向量表征效果，我们自己在网文数据集上继续预训练了一版语言模型，细节参考下文。

Fiction-Bert

句向量特征提取

roberta-tiny预训练结束后，我们将其参数固定，利用语言模型对句子进行向量化编码。客户端输入一段小说内容，返回一个维度为384的向量，该向量能够表征句子段落信息。通过bert-as-service可以方便地将模型部署在服务器上，而且提供了友好的客户端访问方式。

接口调用

推荐方式：

from bert_serving.client import BertClient
bc = BertClient(ip='xx.xx.xx.xx')  # ip address of the GPU machine
bc.encode(['First do it', 'then do it right', 'then do it better'])

同时提供restful api接口：

curl -X POST http://xx.xx.xx.xx:8125/encode \
  -H 'content-type: application/json' \
  -d '{"id": 123,"texts": ["hello world"], "is_tokenized": false}'

注意：若要使用http请求，需要在server端主动开启，否则默认不支持http请求。

工程部署

在对模型进行压缩后，将模型部署在生产的机器上，请求响应的速度还是十分可观的，差不多10ms/it。由于模型比较小巧，占用的显存空间不大，gpu的利用率不高，为了提高单gpu的利用率，我查阅了bert-as-service手册，结合自己的实践（踩坑）经历，总结出几点我认为比较有用的经验。

1. 只要客户端的内存足够，尽量一次encode输入尽可能多的段落，server端会自动对请求的数组进行切分、加入队列，分派给多个worker，worker之间可以并行处理。

2. 根据手册上的benchmark上的数据显示，server端将batch_size调到256后，速度也不会再有更大的提高，gpu利用率也不会提升太多，我自己的实践过程中也是如此。（这点我也没有想明白是为什么，可能是bert-as-service库里实现的问题吧）

3. 若一个gpu上运行一个模型，始终不能达到100%的利用率，可以指定多个worker同时运行在单张gpu上，只要保证显存不要爆就没事，多个worker能够有效地提高整体的运算速度，并且能将gpu的利用率达到100%。

4. 当客户端并发量较大时，遇到了一个坑......当同时向server端发起多个请求时，有时会出现部分任务已经完成了，但是没有返回给客户端，于是就一直hang on着。后来去github上把相关的issue都看了一遍，发现许多人都遇到了和我相似的问题，貌似在版本v1.9.8这个问题被修复了，但是在后面的版本又出现了，我使用的是v1.10.0。最后解决的办法比较粗暴，在客户端那边加一个time out，最好不要用bert-as-service内置的time out，用其他库来做time out限制为上策。

5. max_len参数对模型推理速度以及显存占用影响最大，这很好理解，transformer结构决定了 $O(n^2)$ 时间、空间复杂度。结合具体业务场景，我将这个参数调整为32。

总结与展望

整体项目实现思路其实十分简明，可以简单地概括为：将预训练语言模型作为特征抽取工具，提供向量化服务，在gpu上部署，并提供高性能、可配置、可拓展的服务。

通过这个项目的实践，从采集预训练语料，配置模型参数，到最后预训练模型都撸了一遍，并且针对我们的应用场景，做了一些优化。从模型理解、工程能力、代码实现上都有了些许的提高。

最后要感谢bert4keras、bert-as-service的作者们的开源精神，为NLP在工业落地上起到了很大的推动作用。希望日后自己也能够为开源社区贡献一份微薄力量。