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

绑定权重

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

没有语言模型头的模型没有绑定权重，可以安全导出，而无需 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 实例上。

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

import torch.neuron

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

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

模型架构

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

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

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