TorchScript

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

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

1. 创建虚拟输入以创建模型的*跟踪*，用于序列化到 TorchScript。
2. 对于随机初始化的模型，启用 `~PretrainedConfig.torchscript` 中的 `torchscript` 参数；对于预训练模型，启用 from_pretrained() 中的 `torchscript` 参数。

虚拟输入

虚拟输入用于前向传播，当输入值通过每一层传播时，PyTorch 会跟踪在每个张量上执行的不同操作。记录的操作用于创建模型跟踪。一旦记录，它就会被序列化为 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]

跟踪是根据提供的输入维度创建的，它只能处理与跟踪期间提供的输入具有相同形状的输入。如果输入大小不同，则会显示如下所示的错误消息。

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

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

权重绑定

Transformers 中 `Embedding` 和 `Decoding` 层之间的权重是绑定的，而 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 加载跟踪模型。

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

all_encoder_layers, pooled_output = loaded_model(*dummy_input)

要将跟踪模型用于推理，请使用 `__call__` dunder 方法。

traced_model(tokens_tensor, segments_tensors)

部署到 AWS

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

代替 torch.jit.trace，使用 torch.neuron.trace 来跟踪模型并对其进行优化以适应 Inf1 实例。

import torch.neuron

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

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

模型架构

基于 BERT 的模型——例如 DistilBERT 或 RoBERTa——在 Inf1 实例上运行效果最佳，适用于非生成性任务，例如提取式问答、序列或令牌分类。

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

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