导出到 TorchScript
这是我们开始使用TorchScript进行实验的初期阶段,我们仍在探索其在可变输入大小模型中的能力。这是我们关注的重点,我们将在未来的版本中深化分析,提供更多的代码示例,更灵活的实现,以及比较基于Python的代码与编译后的TorchScript的基准测试。
TorchScript 是一种从 PyTorch 代码创建可序列化和可优化模型的方法。
有两个PyTorch模块,JIT和TRACE,允许开发者导出他们的模型以便在其他程序中重复使用,例如面向效率的C++程序。
我们提供了一个接口,允许你将🤗 Transformers模型导出为TorchScript,以便它们可以在不同于基于PyTorch的Python程序的环境中重复使用。在这里,我们解释了如何使用TorchScript导出和使用我们的模型。
导出模型需要两件事:
- 使用
torchscript
标志进行模型实例化 - 使用虚拟输入进行前向传递
这些必要性意味着开发人员应该注意以下几点,如下所述。
TorchScript 标志和绑定权重
torchscript
标志是必要的,因为大多数 🤗 Transformers 语言模型在其 Embedding
层和 Decoding
层之间有共享权重。TorchScript 不允许导出具有共享权重的模型,因此需要事先解绑并克隆权重。
使用torchscript
标志实例化的模型具有分离的Embedding
层和Decoding
层,这意味着它们不应该在后续进行训练。训练会导致这两层不同步,从而产生意外的结果。
对于没有语言模型头的模型来说,情况并非如此,因为这些模型没有绑定权重。这些模型可以安全地导出,而不需要torchscript
标志。
虚拟输入和标准长度
虚拟输入用于模型的前向传播。当输入值通过各层传播时,PyTorch会跟踪在每个张量上执行的不同操作。这些记录的操作随后用于创建模型的跟踪。
跟踪是根据输入的维度创建的。因此,它受到虚拟输入维度的限制,并且不适用于任何其他序列长度或批量大小。当尝试使用不同的大小时,会引发以下错误:
`The expanded size of the tensor (3) must match the existing size (7) at non-singleton dimension 2`
我们建议您使用至少与推理期间将输入模型的最大输入大小相同的虚拟输入大小来跟踪模型。填充可以帮助填补缺失的值。然而,由于模型是用较大的输入大小进行跟踪的,矩阵的维度也会较大,从而导致更多的计算。
注意每个输入上的操作总数,并在导出不同序列长度模型时密切关注性能。
在 Python 中使用 TorchScript
本节演示如何保存和加载模型,以及如何使用跟踪进行推理。
保存模型
要使用TorchScript导出BertModel
,请从BertConfig
类实例化BertModel
,然后将其保存到磁盘,文件名为traced_bert.pt
:
from transformers import BertModel, BertTokenizer, BertConfig
import torch
enc = BertTokenizer.from_pretrained("google-bert/bert-base-uncased")
# Tokenizing input text
text = "[CLS] Who was Jim Henson ? [SEP] Jim Henson was a puppeteer [SEP]"
tokenized_text = enc.tokenize(text)
# Masking one of the input tokens
masked_index = 8
tokenized_text[masked_index] = "[MASK]"
indexed_tokens = enc.convert_tokens_to_ids(tokenized_text)
segments_ids = [0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1]
# Creating a dummy input
tokens_tensor = torch.tensor([indexed_tokens])
segments_tensors = torch.tensor([segments_ids])
dummy_input = [tokens_tensor, segments_tensors]
# Initializing the model with the torchscript flag
# Flag set to True even though it is not necessary as this model does not have an LM Head.
config = BertConfig(
vocab_size_or_config_json_file=32000,
hidden_size=768,
num_hidden_layers=12,
num_attention_heads=12,
intermediate_size=3072,
torchscript=True,
)
# Instantiating the model
model = BertModel(config)
# The model needs to be in evaluation mode
model.eval()
# If you are instantiating the model with *from_pretrained* you can also easily set the TorchScript flag
model = BertModel.from_pretrained("google-bert/bert-base-uncased", torchscript=True)
# Creating the trace
traced_model = torch.jit.trace(model, [tokens_tensor, segments_tensors])
torch.jit.save(traced_model, "traced_bert.pt")
加载模型
现在你可以从磁盘加载之前保存的 BertModel
,traced_bert.pt
,并在之前初始化的 dummy_input
上使用它:
loaded_model = torch.jit.load("traced_bert.pt")
loaded_model.eval()
all_encoder_layers, pooled_output = loaded_model(*dummy_input)
使用追踪模型进行推理
使用跟踪模型进行推理,通过其__call__
方法:
traced_model(tokens_tensor, segments_tensors)
使用Neuron SDK将Hugging Face TorchScript模型部署到AWS
AWS 推出了 Amazon EC2 Inf1 实例系列,用于在云中进行低成本、高性能的机器学习推理。Inf1 实例由 AWS Inferentia 芯片驱动,这是一种定制的硬件加速器,专门用于深度学习推理工作负载。AWS Neuron 是 Inferentia 的 SDK,支持跟踪和优化 transformers 模型以在 Inf1 上部署。Neuron SDK 提供:
- 易于使用的API,只需一行代码更改即可跟踪和优化TorchScript模型,以便在云端进行推理。
- 开箱即用的性能优化,以提高成本效益。
- 支持使用PyTorch或TensorFlow构建的Hugging Face transformers模型。
影响
基于BERT(双向编码器表示从Transformers)架构的Transformers模型,或其变体如distilBERT和roBERTa,在Inf1上运行最佳,适用于非生成任务,如抽取式问答、序列分类和标记分类。然而,根据此AWS Neuron MarianMT教程,文本生成任务仍然可以适应在Inf1上运行。更多关于可以在Inferentia上直接转换的模型的信息可以在Neuron文档的模型架构适配部分找到。
依赖项
使用AWS Neuron转换模型需要一个Neuron SDK环境,该环境在AWS深度学习AMI上已预先配置。
将模型转换为AWS Neuron
使用与在Python中使用TorchScript相同的代码来跟踪BertModel
,以转换为AWS NEURON模型。导入torch.neuron
框架扩展,以通过Python API访问Neuron SDK的组件:
from transformers import BertModel, BertTokenizer, BertConfig
import torch
import torch.neuron
你只需要修改以下行:
- torch.jit.trace(model, [tokens_tensor, segments_tensors])
+ torch.neuron.trace(model, [tokens_tensor, segments_tensors])
这使得Neuron SDK能够跟踪模型并针对Inf1实例进行优化。
要了解更多关于AWS Neuron SDK的功能、工具、示例教程和最新更新,请参阅AWS NeuronSDK文档。
< > Update on GitHub