Transformers 文件

TorchScript

Hugging Face's logo
加入 Hugging Face 社群

並獲得增強的文件體驗

開始使用

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 例項。

AWS Neuron 需要一個 Neuron SDK 環境,該環境在 AWS DLAMI 上預配置。

代替 torch.jit.trace,使用 torch.neuron.trace 來跟蹤模型並對其進行最佳化以適應 Inf1 例項。

import torch.neuron

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

有關更多資訊,請參閱 AWS Neuron 文件。

模型架構

基於 BERT 的模型——例如 DistilBERTRoBERTa——在 Inf1 例項上執行效果最佳,適用於非生成性任務,例如提取式問答、序列或令牌分類。

文字生成可以適應在 Inf1 例項上執行,如 Transformers MarianMT 教程中所示。

請參閱 推理示例/教程 (Inf1) 指南,瞭解更多關於哪些模型可以開箱即用地轉換為在 Inf1 例項上執行的資訊。

< > 在 GitHub 上更新

© . This site is unofficial and not affiliated with Hugging Face, Inc.