用optimum和fastT5解决Seq2Seq模型推理加速问题

一、问题背景

Seq2Seq模型往往存在模型结构过大，推理速度达不到预期，导致模型推理的接口出现超时的现象。

二、多种解决方法以及优劣性

1.使用NVIDIA提供的TensorRT加速推理模型

链接：https://github.com/NVIDIA/TensorRT

• 优点：根据一些经验贴得知，tensorRT要比下面两种方法有更快的推理速度。(本人没有经过实验验证)
• 缺点：模型转化麻烦，需要先将现有的bin模型转化为ONNX模型，再转化为TRT模型。

2.使用Huggingface官方提供的optimum库

链接：https://huggingface.co/docs/optimum/index

• 优点：因为是huggingface构建的库，对于Transformers库的模型更加易于上手，代码更为简单。
• 缺点：推理速度可能不如TRT模型快。（未实验验证）

3.对于T5类别的模型，可以使用fastT5库

链接：https://github.com/Ki6an/fastT5

• 优点：对于T5类别的模型更加友好，如“t5-base”、“mt5”、“mengzi-t5”等。可以一并生成ONNX模型，并执行推理，代码简单。
• 缺点：1.只能用于T5类别模型，比较局限。2.目前只实现了在CPU上推理，GPU推理需要自己去改写内部代码。

三、实现方法

目前我只使用了optimum方法和fastT5方法，TensorRT暂时没有尝试，后面如果使用了会更新。

1.optimum方法以及步骤

（1）安装需要的环境

我安装的版本是optimum==1.16.0，如果担心版本问题可以按照这个版本

pip install optimum==1.16.0

！！！注意：如果需要用GPU推理一定装的是onnxruntime-gpu
！！！注意下面两个包只能装一个，不然代码一定会出错

# gpu版本
pip install onnxruntime-gpu==1.14.1
# cpu版本
pip install onnxruntime==1.14.1

！！！注意：如果你前面安装的onnxruntime-gpu==1.14.1这个版本,onnx一定不能是1.14版本，会存在版本冲突

pip install onnx==1.13.1

pip install transformers

（2）生成onnx模型

安装好上面的依赖之后，使用optimum进行加速推理之前需要将现有模型转化为onnx模型。直接上命令：

# 如果你是seq2seq任务，一定在task传入 text2text-generation-with-past，这样才能进行seq2seq模型的onnx推理
optimum-cli export onnx --model 你的模型地址 \
                    --task 任务类型，见官方文档 \
                    你的onnx模型保存路径

（3）加载onnx模型，并进行推理

from optimum.onnxruntime import ORTModelForSeq2SeqLM
from transformers import AutoTokenizer

model_path = 你的onnx模型保存路径

# 加载onnx模型
# provider是指定推理的设备，'CUDAExecutionProvider'是GPU，CPUExecutionProvider是用CPU推理
model = ORTModelForSeq2SeqLM.from_pretrained(model_path,provider='CUDAExecutionProvider')
# 加载tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_path)

# 输入input文本
input_text = "请输入你需要输入的文本。"
# 将文本转化为token_ids,不需要用GPU的不用接   .to('cuda')
input_tokens = tokenizer(input_text,return_tensor='pt').to('cuda')
# 产生encoder的输出结果
output_tokens = model.generate(**input_tokens)
# 对于encoder的结果用tokenizer进行解码
outputs = tokenizer.decode(output_tokens)
# 打印输出结果
print(outputs)

以上就是使用optimum加速模型推理的全部过程，如有不对，欢迎留言指正。

2.fastT5加速T5类别模型的推理

（1）安装需要的环境

这个环境就比较简单，不存在依赖冲突，直接安装就行

pip install transformers
pip install fastT5

（2）生成onnx模型和推理一步到位

from fastT5 import export_and_get_onnx_model
from transformers import AutoTokenizer


# 用fastT5输出onnx推理模型
model_name_or_path = "你的现有模型路径"
custom_output_path = "你需要生成的onnx模型路径"
# 生成并加载模型路径
model = export_and_get_onnx_model(model_name_or_path,custom_output_path=custom_output_path,quantized=False)

tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
t_input = "你需要输入的文本，这里可以输入需要测试的内容。"
# 将输入文本进行tokenizer编码变成token_ids
token_input = tokenizer(t_input, return_tensors='pt')
# 调用了model的generate方法，num_beams为集束搜索参数，
tokens_inputs = model.generate(input_ids=token_input['input_ids'],
               attention_mask=token_input['attention_mask'],
               num_beams=1)

# 用tokenizer对结果进行解码，还原成文本
output = tokenizer.decode(tokens_inputs.squeeze(), skip_special_tokens=True)
# 打印结果
print(output)

对于T5模型，如果不考虑使用GPU推理，明显fastT5要简单的多，所以T5类别的模型还是更推荐使用fastT5这个工具包。

总结

个人实验感受来说：
推理速度：TensorRT > optimum > fastT5
易上手度：fastT5 > optimum > TensorRT
对于我自己实验的模型来说，fastT5要比optimum实现的模型推理速度稍慢一点，其他模型我没有进行验证。
以上内容纯属个人观点，如果有写的不准确的地方，欢迎大家及时指正。