将 T5 模型大小缩减至原来的三分之一，并将推理速度提升至最多五倍。

• 安装
• 使用方法
• 详细信息
• 功能特性
• 基准测试
  • ONNX 模型
  • 量化 ONNX 模型
• 量化模型评分
• 进一步改进
• 许可证
• 获取帮助
• 致谢

T5 模型可用于多种自然语言处理任务，如摘要生成、问答、问题生成、翻译、文本生成等。由于序列式文本生成的特性，其推理过程通常较为缓慢，而较大的 T5 模型则会更加耗时。fastT5 通过在 ONNX Runtime 上运行 T5 模型，显著提升了推理速度；同时，通过对模型进行量化，进一步减小了模型体积。

fastT5 库允许用户将预训练的 T5 模型转换为 ONNX 格式，对其进行量化，并以单行代码的形式输出可在 ONNX Runtime 上运行的模型。此外，用户还可以自定义整个流程。

---

安装

您可以通过 PyPI 安装 fastT5：

 pip install fastt5

如果您希望从源码构建：

git clone https://github.com/Ki6an/fastT5
cd fastT5
pip3 install -e .

使用方法

export_and_get_onnx_model() 方法会将给定的预训练 T5 模型导出为 ONNX 格式，对其进行量化，并使用默认设置在 ONNX Runtime 上运行。该方法返回的模型支持 Hugging Face 的 generate() 方法。

如果您不希望对模型进行量化，可以在方法中设置 quantized=False。

from fastT5 import export_and_get_onnx_model
from transformers import AutoTokenizer

model_name = 't5-small'
model = export_and_get_onnx_model(model_name)

tokenizer = AutoTokenizer.from_pretrained(model_name)
t_input = "translate English to French: The universe is a dark forest."
token = tokenizer(t_input, return_tensors='pt')

tokens = model.generate(input_ids=token['input_ids'],
               attention_mask=token['attention_mask'],
               num_beams=2)

output = tokenizer.decode(tokens.squeeze(), skip_special_tokens=True)
print(output)

若要运行已导出的模型，请使用 get_onnx_model()。

您可以按照以下代码示例自定义整个流程：

from fastT5 import (OnnxT5, get_onnx_runtime_sessions,
                    generate_onnx_representation, quantize)
from transformers import AutoTokenizer

model_or_model_path = 't5-small'

# 步骤 1：将 Hugging Face 的 T5 模型转换为 ONNX 格式
onnx_model_paths = generate_onnx_representation(model_or_model_path)

# 步骤 2：（推荐）对转换后的模型进行量化，以加快推理速度并减小模型体积。
quant_model_paths = quantize(onnx_model_paths)

# 步骤 3：设置 ONNX Runtime
model_sessions = get_onnx_runtime_sessions(quant_model_paths)

# 步骤 4：获取 ONNX 模型
model = OnnxT5(model_or_model_path, model_sessions)

                      ...

自定义输出路径

默认情况下，fastT5 会在当前目录下创建一个 models 文件夹，并将所有模型存储其中。您可以指定一个自定义文件夹路径来存储导出的模型。若要运行已存储在自定义路径下的导出模型，请使用 get_onnx_model(onnx_models_path="/path/to/custom/folder/")。

from fastT5 import export_and_get_onnx_model, get_onnx_model

model_name = "t5-small"
custom_output_path = "/path/to/custom/folder/"

# 1. 将模型存储到自定义输出路径
model = export_and_get_onnx_model(model_name, custom_output_path)

# 2. 运行已存储在自定义路径下的导出模型
# model = get_onnx_model(model_name, custom_output_path)

详细信息

T5 是一种序列到序列（Encoder-Decoder）模型。由于其在推理过程中会反复使用解码器，因此我们无法直接将整个模型导出为 ONNX 格式。我们需要分别导出编码器和解码器。

past_key_values 包含预先计算的隐藏状态（自注意力层和交叉注意力层中的键值），可用于加速序列解码。

模型只能以固定数量的输入进行导出。然而，第一步的解码器并不需要 past_key_values，而后续步骤的解码器则需要。为了解决这一问题，我们可以创建两个解码器：一个用于第一步，不使用 past_key_values；另一个用于后续步骤，利用 past_key_values。

接下来，我们将导出三个模型（编码器、解码器、初始解码器）。然后对它们进行量化，将 32 位量化为 8 位，理论上可使内存占用减少四分之一。由于多了一个解码器，模型体积将减少约三分之二。

最后，我们将量化后的模型部署到 ONNX Runtime 上运行。

推理过程非常简单，因为该模型支持 Hugging Face 的 generate() 方法。

功能特性

• 轻松将任何预训练的 T5 模型导出为 ONNX 格式（包含 past_key_values）。
• 导出的模型通过 generate() 方法支持束搜索、贪婪搜索等多种解码策略。
• 通过量化技术将模型体积缩小至原来的三分之一。
• 相比 PyTorch 执行，贪婪搜索的速度可提升至五倍，束搜索的速度可提升至三到四倍。

基准测试

以下基准测试结果基于 T5-base 模型在英法翻译任务上的表现。

ONNX 模型

下图展示了量化后的 ONNX 模型与 PyTorch 模型在束搜索数量从 1 到 9 不同时的延迟情况。此处显示的延迟是针对序列长度不超过 130 的平均值。

以下热力图显示了 PyTorch 模型与 ONNX 模型延迟的倍数关系。大多数情况下，ONNX 模型的表现更优。然而，当序列长度较长时，ONNX 模型的速度会有所下降。

量化后的 ONNX 模型

如前所述，量化模型是轻量级模型，其准确率几乎与原始模型相同（量化模型的评分将在下一节中提及）。与 ONNX 和 PyTorch 模型相比，量化后的 ONNX 模型具有最低的延迟。

该模型在贪婪搜索上的平均性能比 PyTorch 模型快 5.7 倍，在束搜索上的性能则快 3 到 4 倍。

注意：这些结果是在 AMD EPYC 7B12 上生成的，不同设备的结果可能会有所差异。ONNX 模型通常在核心数较多的高端 CPU 上表现良好。

量化模型评分

测试结果为英语到法语的翻译任务，使用束搜索，束宽设为 3。

	Bleu_4	METEOR	ROUGE_L
t5-small（量化）	0.240769	0.282342	0.468817
t5-small（PyTorch）	0.254601	0.295172	0.492749
t5-base（量化）	0.267606	0.306019	0.499188
t5-base（PyTorch）	0.268346	0.304969	0.503306
t5-large（量化）	0.286726	0.316845	0.503585
t5-large（PyTorch）	0.294015	0.315774	0.508677

私有 HuggingFace Model Hub 模型

HuggingFace Model Hub 支持私有模型。要使用 fastT5 的私有预训练 T5 版本，您首先需要通过 $ transformers-cli login 在 HuggingFace 生态系统中进行身份验证。然后，在使用 fastT5 时，还需要额外的导入和调用：

from fastT5 import (
    OnnxT5,
    get_onnx_runtime_sessions,
    generate_onnx_representation,
    quantize,
    set_auth_token)
from transformers import AutoTokenizer

set_auth_token(True)
# 其余代码与使用公开模型相同

如果您无法运行 $ transformers-cli login，或者更倾向于使用您的 API 密钥（可在 https://huggingface.co/settings/token 或组织的 https://huggingface.co/organizations/ORG_NAME/settings/token 页面找到），您可以将该密钥作为字符串传递给 set_auth_token。为了避免将 API 密钥硬编码到代码中，可以设置环境变量 HF_API_KEY=<redacted>，然后在代码中：

import os

from fastT5 import (
    OnnxT5,
    get_onnx_runtime_sessions,
    generate_onnx_representation,
    quantize,
    set_auth_token)
from transformers import AutoTokenizer

auth_token = os.environ.get("HF_API_KEY")
set_auth_token(auth_token)

# 代码继续正常执行

未来改进方向

• 目前 fastT5 库仅支持 ONNX Runtime 的 CPU 版本，GPU 实现仍有待完成。
• 对 ONNX 模型进行图优化将进一步降低延迟。

获取帮助

• 请联系我：kiranr8k@gmail.com
• 如有需要，请在 GitHub 上 提交问题。

致谢

• 原始 T5 论文
• HuggingFace 的 transformers
• ONNX
• Microsoft 的 ONNX Runtime
• onnxt5

@article{2019t5,
  author = {Colin Raffel、Noam Shazeer、Adam Roberts、Katherine Lee、Sharan Narang、Michael Matena、Yanqi Zhou、Wei Li、Peter J. Liu},
  title = {利用统一的文本到文本 Transformer 探索迁移学习的极限},
  journal = {arXiv e-prints},
  year = {2019},
  archivePrefix = {arXiv},
  eprint = {1910.10683},
}

fastT5 快速上手指南

fastT5 是一个用于加速 Hugging Face T5 模型推理的开源库。它通过将模型转换为 ONNX 格式并进行量化，可将模型体积缩小 3 倍，推理速度提升最高 5 倍，同时保持与原始模型相近的精度。

环境准备

• 操作系统: Linux, macOS, Windows
• Python 版本: 建议 Python 3.7+
• 核心依赖:
  • transformers (Hugging Face)
  • onnx
  • onnxruntime (当前版本主要支持 CPU 推理)

安装步骤

你可以直接通过 PyPI 安装 fastT5：

pip install fastt5

如果需要从源码构建：

git clone https://github.com/Ki6an/fastT5
cd fastT5
pip3 install -e .

国内加速提示：如果下载速度较慢，建议使用国内镜像源安装：

pip install fastt5 -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

fastT5 提供了最简化的接口 export_and_get_onnx_model()，只需一行代码即可完成模型导出、量化并加载到 ONNX Runtime 中。返回的模型完全兼容 Hugging Face 的 generate() 方法。

以下是一个完整的翻译示例（英语转法语）：

from fastT5 import export_and_get_onnx_model
from transformers import AutoTokenizer

# 1. 加载并转换模型 (自动导出为 ONNX 并量化)
model_name = 't5-small'
model = export_and_get_onnx_model(model_name)

# 2. 加载分词器
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 3. 准备输入数据
t_input = "translate English to French: The universe is a dark forest."
token = tokenizer(t_input, return_tensors='pt')

# 4. 生成结果 (用法与原生 PyTorch 模型一致)
tokens = model.generate(
    input_ids=token['input_ids'],
    attention_mask=token['attention_mask'],
    num_beams=2
)

# 5. 解码输出
output = tokenizer.decode(tokens.squeeze(), skip_special_tokens=True)
print(output)

自定义存储路径： 默认情况下，转换后的模型会保存在当前目录的 models 文件夹中。你可以指定自定义路径：

custom_output_path = "/path/to/custom/folder/"

# 将模型导出并保存到指定路径
model = export_and_get_onnx_model(model_name, custom_output_path)

# 或者直接加载已存在于指定路径的模型
# model = get_onnx_model(model_name, custom_output_path)