Transformers 文档

TorchScript

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

TorchScript

TorchScript 将 PyTorch 模型序列化为可以在非 Python 进程中执行的程序。这在生产环境中尤其有利,因为在这些环境中,Python 可能不是性能最高的选择。

Transformers 可以通过以下方式将模型导出到 TorchScript:

  1. 创建虚拟输入以创建模型的trace,从而序列化到 TorchScript
  2. ~PretrainedConfig.torchscript 中为随机初始化的模型或为预训练模型启用 torchscript 参数 from_pretrained()

虚拟输入

虚拟输入在正向传递中使用,并且当输入值通过每一层传播时,PyTorch 会跟踪在每个张量上执行的不同操作。记录的操作用于创建模型 trace。一旦记录下来,它就会被序列化为 TorchScript 程序。

from transformers import BertModel, BertTokenizer, BertConfig
import torch

tokenizer = BertTokenizer.from_pretrained("google-bert/bert-base-uncased")
text = "[CLS] Who was Jim Henson ? [SEP] Jim Henson was a puppeteer [SEP]"
tokenized_text = tokenizer.tokenize(text)

masked_index = 8
tokenized_text[masked_index] = "[MASK]"
indexed_tokens = tokenizer.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]

trace 是基于提供的输入维度创建的,并且在 tracing 期间它只能处理与提供的输入形状相同的输入。不同大小的输入会引发如下所示的错误消息。

`The expanded size of the tensor (3) must match the existing size (7) at non-singleton dimension 2`.

尝试创建一个 trace,其虚拟输入大小至少与推理期间的最大预期输入一样大。填充可以帮助填充较大输入的缺失值。但是,由于较大的输入大小需要更多的计算,因此可能会较慢。请注意在每个输入上执行的操作总数,并在导出具有可变序列长度的模型时跟踪模型性能。

绑定权重

EmbeddingDecoding 层之间的权重在 Transformers 中是绑定的,而 TorchScript 无法导出具有绑定权重的模型。使用 torchscript=True 实例化模型会分离 EmbeddingDecoding 层,并且它们不会再接受任何训练,因为这会将两层置于不同步状态,从而可能导致意外结果。

没有语言模型头的模型没有绑定权重,可以安全导出,而无需 torchscript 参数。

随机初始化的模型
预训练模型
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,
)

model = BertModel(config)
model.eval()

导出到 TorchScript

使用 torch.jit.trace 创建 TorchScript 程序,并使用 torch.jit.save 保存。

traced_model = torch.jit.trace(model, [tokens_tensor, segments_tensors])
torch.jit.save(traced_model, "traced_bert.pt")

使用 torch.jit.load 加载 traced 模型。

loaded_model = torch.jit.load("traced_bert.pt")
loaded_model.eval()

all_encoder_layers, pooled_output = loaded_model(*dummy_input)

要使用 traced 模型进行推理,请使用 __call__ dunder 方法。

traced_model(tokens_tensor, segments_tensors)

部署到 AWS

从 Transformers 序列化的 TorchScript 程序可以部署在 Amazon EC2 Inf1 实例上。该实例由 AWS Inferentia 芯片提供支持,这是一种为深度学习推理工作负载设计的自定义硬件加速器。AWS Neuron 支持 tracing Transformers 模型以部署在 Inf1 实例上。

AWS Neuron 需要 Neuron SDK 环境,该环境在 AWS DLAMI 上预配置。

使用 torch.neuron.trace 而不是 torch.jit.trace 来 trace 模型并针对 Inf1 实例对其进行优化。

import torch.neuron

torch.neuron.trace(model, [tokens_tensor, segments_tensors])

有关更多信息,请参阅 AWS Neuron 文档。

模型架构

基于 BERT 的模型(如 DistilBERTRoBERTa)在 Inf1 实例上运行效果最佳,适用于非生成任务,例如抽取式问答以及序列或 token 分类。

文本生成可以调整为在 Inf1 实例上运行,如 Transformers MarianMT 教程中所示。

有关哪些模型可以开箱即用地转换为在 Inf1 实例上运行的更多信息,请参阅 推理示例/教程 (Inf1) 指南。

< > 在 GitHub 上更新