1. 磐创AI首页
  2. 系列教程
  3. Transformers

Transformers 中使用 TorchScript | 四

本文是全系列中第6 / 8篇:Transformers

作者|huggingface
编译|VK
来源|Github

注意:这是我们使用TorchScript进行实验的开始,我们仍在探索可变输入大小模型的功能。它是我们关注的焦点,我们将在即将发布的版本中加深我们的分析,提供更多代码示例,更灵活的实现以及将基于python的代码与已编译的TorchScript进行基准测试的比较。

根据Pytorch的文档:”TorchScript是一种从PyTorch代码创建可序列化和可优化模型的方法”。Pytorch的两个模块JIT和TRACE允许开发人员导出他们的模型,这些模型可以在其他程序中重用,例如面向效率的C++程序。

我们提供了一个接口,该接口允许将transformers模型导出到TorchScript,以便他们可在与基于Pytorch的python程序不同的环境中重用。在这里,我们解释如何使用我们的模型,以便可以导出它们,以及将这些模型与TorchScript一起使用时要注意的事项。

导出模型需要两件事:
– 虚拟化输入以执行模型正向传播。
– 需要使用torchscript标志实例化该模型。

这些必要性意味着开发人员应注意几件事。这些在下面详细说明。

含义

TorchScript标志和解绑权重

该标志是必需的,因为该存储库中的大多数语言模型都在它们的Embedding层及其Decoding层具有绑定权重关系。TorchScript不允许导出绑定权重的模型,因此,有必要事先解绑权重。

这意味着以torchscript标志实例化的模型使得Embedding层和Decoding层分开,这意味着不应该对他们进行同时训练,导致意外的结果。

对于没有语言模型头(Language Model head)的模型,情况并非如此,因为那些模型没有绑定权重。这些型号可以在没有torchscript标志的情况下安全地导出。

虚拟(dummy)输入和标准长度

虚拟输入用于进行模型前向传播。当输入值在各层中传播时,Pytorch跟踪在每个张量上执行的不同操作。然后使用这些记录的操作创建模型的“迹”。

迹是相对于输入的尺寸创建的。因此,它受到虚拟输入尺寸的限制,并且不适用于任何其他序列长度或批次大小。尝试使用其他尺寸时,会出现如下错误,如:

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

因此,建议使用至少与最大输入大小相同的虚拟输入大小来跟踪模型。在推理期间对于模型的输入,可以执行填充来填充缺少的值。作为模型
将以较大的输入大小来进行跟踪,但是,不同矩阵的尺寸也将很大,从而导致更多的计算。

建议注意每个输入上完成的操作总数,并密切关注各种序列长度对应性能的变化。

在Python中使用TorchScript

以下是使用Python保存,加载模型以及如何使用”迹”进行推理的示例。

保存模型

该代码段显示了如何使用TorchScript导出BertModel。在这里实例化BertModel,根据BertConfig类,然后以文件名traced_bert.pt保存到磁盘

from transformers import BertModel, BertTokenizer, BertConfig
import torch

enc = BertTokenizer.from_pretrained("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标志初始化模型
# 标志被设置为True,即使没有必要,因为该型号没有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)

# 实例化模型
model = BertModel(config)

# 模型设置为评估模式
model.eval()

# 如果您要​​使用from_pretrained实例化模型,则还可以设置TorchScript标志
model = BertModel.from_pretrained("bert-base-uncased", torchscript=True)

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

载入模型

该代码段显示了如何加载以前以名称traced_bert.pt保存到磁盘的BertModel
我们重新使用之前初始化的dummy_input

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

all_encoder_layers, pooled_output = loaded_model(dummy_input)

使用跟踪模型进行推理

使用跟踪模型进行推理就像使用其__call__ 方法一样简单:

traced_model(tokens_tensor, segments_tensors)

原文链接:https://huggingface.co/transformers/torchscript.html

原创文章,作者:磐石,如若转载,请注明出处:https://panchuang.net/2020/03/25/transformers-%e4%b8%ad%e4%bd%bf%e7%94%a8-torchscript-%e5%9b%9b/

发表评论

登录后才能评论

联系我们

400-800-8888

在线咨询:点击这里给我发消息

邮件:admin@example.com

工作时间:周一至周五,9:30-18:30,节假日休息