导出到 TorchScript¶
这是我们对 TorchScript 进行实验的起点，我们仍在探索它在处理可变输入大小模型方面的能力。这是我们感兴趣的重点，未来我们将进一步深入分析，并提供更多代码示例、更灵活的实现以及与基于 Python 的代码相比的性能基准测试。
根据 TorchScript 文档：
TorchScript 是一种从 PyTorch 代码创建可序列化和可优化模型的方法。
有两个 PyTorch 模块，JIT 和 TRACE，允许开发者导出他们的模型以供其他程序（如效率导向的 C++ 程序）重用。
我们提供了一个接口，允许您将 🤗 Transformers 模型导出到 TorchScript，以便在不同的环境中重用这些模型，而不仅仅是基于 PyTorch 的 Python 程序。在这里，我们将解释如何导出和使用这些模型。
导出模型需要两件事：
- 使用
torchscript标志实例化模型 - 使用虚拟输入进行前向传递
这些要求意味着开发者需要注意以下几点。
使用 torchscript 标志和绑定权重¶
torchscript 标志是必需的，因为大多数 🤗 Transformers 语言模型在其 Embedding 层和 Decoding 层之间有绑定权重。TorchScript 不允许导出具有绑定权重的模型，因此需要预先解开并克隆这些权重。
使用 torchscript 标志实例化的模型，其 Embedding 层和 Decoding 层是分开的，这意味着它们不应再进行训练。继续训练会导致这两个层不同步，从而产生意外结果。
对于没有语言模型头的模型，情况并非如此，因为这些模型没有绑定权重。这些模型可以安全地导出，而无需使用 torchscript 标志。
虚拟输入和标准长度¶
虚拟输入用于模型的前向传递。当输入值通过各层传播时，PyTorch 会记录每个张量上执行的不同操作。这些记录的操作用于创建模型的 trace。
trace 是相对于输入维度创建的，因此受虚拟输入维度的限制，对于任何其他序列长度或批量大小都不起作用。尝试使用不同大小的输入时，会引发以下错误：
from transformers import BertModel, BertTokenizer, BertConfig
import torch
# 加载预训练的分词器
enc = BertTokenizer.from_pretrained("google-bert/bert-base-uncased")
# 对输入文本进行分词
text = "[CLS] Who was Jim Henson ? [SEP] Jim Henson was a puppeteer [SEP]"
tokenized_text = enc.tokenize(text)
# 遮蔽一个输入词
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]
# 创建虚拟输入
tokens_tensor = torch.tensor([indexed_tokens])
segments_tensors = torch.tensor([segments_ids])
dummy_input = [tokens_tensor, segments_tensors]
# 使用 torchscript 标志初始化模型
# 尽管这个模型没有 LM 头，但仍然设置为 True
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()
# 如果使用 *from_pretrained* 实例化模型，也可以轻松设置 TorchScript 标志
model = BertModel.from_pretrained("google-bert/bert-base-uncased", torchscript=True)
# 创建跟踪
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 是支持在 Inf1 上部署 Transformer 模型的 SDK。Neuron SDK 提供以下功能：
- 使用一行代码更改即可轻松跟踪和优化 TorchScript 模型以在云端进行推理的易用 API。
- 开箱即用的性能优化，以提高 成本效益。
- 支持使用 PyTorch 或 TensorFlow 构建的 Hugging Face Transformer 模型。
影响¶
基于 BERT (双向编码器表示来自 Transformer) 架构或其变体（如 distilBERT 和 roBERTa）的 Transformer 模型在 Inf1 上最适合非生成任务，例如提取性问答、序列分类和 token 分类。然而，文本生成任务仍可以根据此 AWS Neuron MarianMT 教程 适应在 Inf1 上运行。更多关于可以直接在 Inferentia 上转换的模型的信息，请参阅 Neuron 文档中的 模型架构适配 部分。
依赖项¶
使用 AWS Neuron 转换模型需要一个 Neuron SDK 环境，该环境预先配置在 AWS Deep Learning AMI 上。
将模型转换为 AWS Neuron¶
使用与 在 Python 中使用 TorchScript 相同的代码跟踪 BertModel。导入 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 文档。