官网链接
Fast Transformer Inference with Better Transformer — PyTorch Tutorials 2.0.1+cu117 documentation
使用 BETTER TRANSFORMER 快速的推理TRANSFORMER
本教程介绍了作为PyTorch 1.12版本的一部分的Better Transformer (BT)。在本教程中，我们将展示如何使用更好的Transformer与torchtext进行生产推理。Better Transformer是一个具备生产条件fastpath并且可以加速在CPU和GPU上具有高性能的Transformer模型的部署。对于直接基于PyTorch核心nn.module或基于torchtext的模型，fastpath功能可以透明地工作。
使用PyTorch核心torch.nn.module类TransformerEncoder, TransformerEncoderLayer和MultiHeadAttention的模型，可以通过Better Transformer fastpath 执行加速。此外，torchtext已经更新为使用核心库模块，以受益于fastpath加速。(将来可能会启用其他模块的fastpath执行。)
Better Transformer提供两种类型的加速:
- 实现CPU和GPU的Native multihead attention(MHA)，提高整体执行效率。
- 利用NLP推理中的稀疏性。由于输入长度可变，输入令牌可能包含大量填充令牌，可以跳过处理，从而显著提高速度。
Fastpath执行受制于一些标准。最重要的是，模型必须在推理模式下执行，并且在不收集梯度信息的输入张量上运行(例如，使用torch.no_grad运行)。
本教程中Better Transformer 特点
- 加载预训练模型(1.12之前没有Better Transformer)
- 在CPU上并且没有BT fastpath(仅本机MHA))的情况下 运行和基准测试推断
- 在设备(可配置)上并且没有BT fastpath(仅本机MHA))的情况下 运行和基准测试推断
- 启用稀疏性支持
- 在设备(可配置)上并且没有BT fastpath(仅本机MHA+稀疏性))的情况下 运行和基准测试推断
额外的信息
关于Better Transformer的其他信息可以在PyTorch.Org 博客中找到。A Better Transformer for Fast Transformer Inference.
设置
加载预训练模型
我们按照torchtext.models中的说明从预定义的torchtext模型下载XLM-R模型。我们还将DEVICE设置为执行加速器上的测试。(根据您的环境适当启用GPU执行。)
import torch
import torch.nn as nnprint(f"torch version: {torch.__version__}")DEVICE = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")print(f"torch cuda available: {torch.cuda.is_available()}")import torch, torchtext
from torchtext.models import RobertaClassificationHead
from torchtext.functional import to_tensor
xlmr_large = torchtext.models.XLMR_LARGE_ENCODER
classifier_head = torchtext.models.RobertaClassificationHead(num_classes=2, input_dim = 1024)
model = xlmr_large.get_model(head=classifier_head)
transform = xlmr_large.transform()数据集搭建
我们设置了两种类型的输入:一个小的输入批次和一个具有稀疏性的大的输入批次。
small_input_batch = ["Hello world","How are you!"
]
big_input_batch = ["Hello world","How are you!","""`Well, Prince, so Genoa and Lucca are now just family estates of the
Buonapartes. But I warn you, if you don't tell me that this means war,
if you still try to defend the infamies and horrors perpetrated by
that Antichrist- I really believe he is Antichrist- I will have
nothing more to do with you and you are no longer my friend, no longer
my 'faithful slave,' as you call yourself! But how do you do? I see
I have frightened you- sit down and tell me all the news.`It was in July, 1805, and the speaker was the well-known Anna
Pavlovna Scherer, maid of honor and favorite of the Empress Marya
Fedorovna. With these words she greeted Prince Vasili Kuragin, a man
of high rank and importance, who was the first to arrive at her
reception. Anna Pavlovna had had a cough for some days. She was, as
she said, suffering from la grippe; grippe being then a new word in
St. Petersburg, used only by the elite."""
]接下来，我们选择小批量或大批量输入，对输入进行预处理并测试模型。
input_batch=big_input_batchmodel_input = to_tensor(transform(input_batch), padding_value=1)
output = model(model_input)
output.shape最后，我们设置基准迭代计数:
ITERATIONS=10执行
在CPU上并且没有BT fastpath(仅本机MHA)的情况下 运行和基准测试推断
我们在CPU上运行模型，并收集概要信息:
- 第一次运行使用传统方式(“slow path”)执行。
- 第二次运行通过使用model.eval()将模型置于推理模式来启用BT fastpath执行，并使用torch.no_grad()禁用梯度收集。
当模型在CPU上执行时，您可以看到改进(其大小取决于CPU模型)。注意，fastpath配置文件显示了本机TransformerEncoderLayer实现aten::_transformer_encoder_layer_fwd.中的大部分执行时间。
print("slow path:")
print("==========")
with torch.autograd.profiler.profile(use_cuda=False) as prof:for i in range(ITERATIONS):output = model(model_input)
print(prof)model.eval()print("fast path:")
print("==========")
with torch.autograd.profiler.profile(use_cuda=False) as prof:with torch.no_grad():for i in range(ITERATIONS):output = model(model_input)
print(prof)
在设备(可配置)上并且没有BT fastpath(仅本机MHA))的情况下 运行和基准测试推断
我们检查BT 稀疏性设置:
model.encoder.transformer.layers.enable_nested_tensor我们禁用BT 稀疏性:
model.encoder.transformer.layers.enable_nested_tensor=False
我们在DEVICE上运行模型，并收集DEVICE上本机MHA执行的配置文件信息:
- 第一次运行使用传统方式(“slow path”)执行。
- 第二次运行通过使用model.eval()将模型置于推理模式来启用BT fastpath执行，并使用torch.no_grad()禁用梯度收集。
当在GPU上执行时，你应该看到一个显著的加速，特别是对于包含稀疏性的大输入批处理设置:
model.to(DEVICE)
model_input = model_input.to(DEVICE)print("slow path:")
print("==========")
with torch.autograd.profiler.profile(use_cuda=True) as prof:for i in range(ITERATIONS):output = model(model_input)
print(prof)model.eval()print("fast path:")
print("==========")
with torch.autograd.profiler.profile(use_cuda=True) as prof:with torch.no_grad():for i in range(ITERATIONS):output = model(model_input)
print(prof)总结
在本教程中，我们介绍了使用 Better Transformer fastpath快速的transformer 推理，在torchtext 中使用PyTorch核心的 Better Transformer包支持Transformer Encoder 模型。在确认BT fastpath可用性的前提下，我们已经演示了 Better Transformer 的使用。我们已经演示并测试了BT fastpath 执行模式·、本机MHA执行和BT稀疏性加速的使用。
