9n-triton部署bert模型实战经验
一、背景
对于算法工程师来说,通常采用python语言来作为工作语言,但是直接用python部署线上服务性能很差。这个问题困扰了我很久,为了缓解深度学习模型工程落地性能问题,探索了Nvidia提供的triton部署框架,并在九数中台上完成线上部署,发现性能提升近337%!!(原服务单次访问模型推理时间175ms左右,同模型框架单次访问推理时间缩减至40ms左右)。本文仅以本次部署服务经验分享给诸多受制于python性能问题的伙伴,希望有所帮助。
二、问题发现
组内原算法服务一直采用python-backend的镜像部署方式,将算法模型包装成接口形式,再通过Flask外露,打入docker中启动服务,但是发现推到线上接口响应时间过长,非常影响用户体验,于是想做出改进。python后端部署一般存在以下问题:
为了解决上述困扰,我们经过调研行业内其他实践经验,决定摒弃传统镜像部署方式,摆脱python部署环境,换用triton部署算法模型。
三、部署实例
本试验通过fine-tune的Bert模型进行文本分类任务,并通过九数算法中台包装的triton部署框架部署至公网可访问,以下介绍了我们团队部署的全流程和途中踩到的坑,以供大家快速入手(以下均以Bert模型为例):
step 1: 将训练好的模型保存为onnx或torchscript(即.pt)文件格式。据调研行业经验,onnx一般用于cpu环境下推理任务效果较好,.pt文件在gpu环境下推理效果较好,我们最终选择了将模型转化为.pt文件。代码例子如下:
import torch
from transformers import BertTokenizer, BertForSequenceClassification
tokenizer = BertTokenizer.from_pretrained('./pretrained-weights/bert-base-chinese')
model = BertForSequenceClassification.from_pretrained('./pretrained-weights/bert-base-chinese', num_labels=10)
# Trace the wrapped model
# input_ids和attention_mask来自tokenizer, 具体可看一下torch.jit.trace用法
traced_model = torch.jit.trace(model, (input_ids, attention_mask))
traced_model.save("./saved_model_torchscript_v3.pt")
本段代码借用了torch.jit.trace的存储方式,输入是文本经过tokenizer之后的input_ids和attention_mask,输出是未经softmax处理的线性层。我们尝试将softmax部分包装进模型中一起打包成pt文件,使模型能够直接吐出分类类别,方案是外层再包装一层forward,代码例子如下:
class BertSequenceClassifierWrapper(torch.nn.Module):
def __init__(self, model):
super().__init__()
self.model = model.cuda()
def forward(self, input_ids, attention_mask):
outputs = self.model(input_ids=input_ids.cuda(), attention_mask=attention_mask.cuda())
# Apply softmax to the logits
probs = softmax(outputs.logits, dim=-1)
# Get the predicted class
predicted_class = torch.argmax(probs, dim=-1)
return predicted_class
wrapped_model = BertSequenceClassifierWrapper(model)
traced_model = torch.jit.trace(wrapped_model,