TurboTransformers 加速预训练模型推理

TurboTransformers 是腾讯开源的BERT推理模型
TurboTransformers 支持 python 和 C++接口进行调用。TurboTransformers 支持 TensorFlow 和 PyTorch 预训练模型的载入。它可以作为 huggingface/transformers 的推理加速插件,通过加入几行 python 代码获得 BERT 模型的端对端加速效果。
和 ONNX-runtime、TensorRT、Torchlib 等推理优化引擎相比,TurboTransformers 在性能和使用方式上都具备优势。

TurboTransformers 加速预训练模型推理_第1张图片


使用方式:
1. 安装TurboTansformers容器环境
1.1 下载镜像
可参考官方github地址上的步骤进行构建,也可以直接下载官方已经构造好的cpu/gpu镜像(我是这么操作的)
docker pull thufeifeibear/turbo_transformers_cpu:latest
1.2 构造自己的容器
容器使用,参考菜鸟教程:Docker 容器使用 | 菜鸟教程
1.3 进入容器,安装自己代码依赖的python包
docker exec -it turbo_transformer /bin/bash
1.4 上传本地代码到容器中
例: docker cp /Users/fangbing/code/seo_keyword 2ad2439e6cc4:/home
2. 代码改写
参考官方github地址的示例进行改写,也可查看源码,挑选合适的模型方法,目前支持的模型结构如下:

TurboTransformers 加速预训练模型推理_第2张图片


使用示例:
使用sentence-transformers预训练模型进行相似度计算,直接上线时间性能上不满足要求,可使用TurboTransformers 进行推理加速
# TuboTransformers在CPU下默认使用的是ONNX,在GPU下才使用tubo

TurboTransformers 加速预训练模型推理_第3张图片


使用也比较简单:

tokenizer = AutoTokenizer.from_pretrained("./all-MiniLM-L6-v2")
model = AutoModel.from_pretrained("./all-MiniLM-L6-v2")
encoded_input = tokenizer(sentences, padding=True, truncation=True, max_length=128, return_tensors='pt')

tt_model = turbo_transformers.BertModel.from_torch(model)

res=tt_model(encoded_input['input_ids'],token_type_ids=encoded_input['token_type_ids'],attention_masks=encoded_input['attention_mask'])

你可能感兴趣的:(自然语言处理,python,深度学习)