#
Export to TorchScript

これはTorchScriptを使用した実験の最初であり、可変入力サイズのモデルに対するその能力をまだ探求中です。これは私たちの関心の焦点であり、今後のリリースでは、より柔軟な実装や、PythonベースのコードとコンパイルされたTorchScriptを比較するベンチマークを含む、より多くのコード例で詳細な分析を行います。

TorchScriptのドキュメントによれば：

> TorchScriptは、PyTorchコードから直列化および最適化可能なモデルを作成する方法です。

TorchScriptを使用すると、効率志向のC++プログラムなど、他のプログラムでモデルを再利用できるようになります。PyTorchベースのPythonプログラム以外の環境で🤗 Transformersモデルをエクスポートして使用するためのインターフェースを提供しています。ここでは、TorchScriptを使用してモデルをエクスポートし、使用する方法を説明します。

モデルをエクスポートするには、次の2つの要件があります：

- `torchscript`フラグを使用したモデルのインスタンス化
- ダミーの入力を使用したフォワードパス

これらの必要条件は、以下で詳細に説明されているように、開発者が注意する必要があるいくつかのことを意味します。

##
TorchScript flag and tied weights

`torchscript`フラグは、ほとんどの🤗 Transformers言語モデルにおいて、`Embedding`レイヤーと`Decoding`レイヤー間で重みが連結されているため必要です。 TorchScriptでは、重みが連結されているモデルをエクスポートすることはできませんので、事前に重みを切り離して複製する必要があります。

`torchscript`フラグを使用してインスタンス化されたモデルは、`Embedding`レイヤーと`Decoding`レイヤーが分離されており、そのため後でトレーニングしてはいけません。 トレーニングは、これらの2つのレイヤーを非同期にする可能性があり、予期しない結果をもたらす可能性があります。

言語モデルヘッドを持たないモデルには言及しませんが、これらのモデルには連結された重みが存在しないため、`torchscript`フラグなしで安全にエクスポートできます。

##
Dummy inputs and standard lengths

ダミー入力はモデルのフォワードパスに使用されます。入力の値はレイヤーを通じて伝播される間、PyTorchは各テンソルに実行された異なる操作を追跡します。これらの記録された操作は、モデルのトレースを作成するために使用されます。

トレースは入力の寸法に対して作成されます。そのため、ダミー入力の寸法に制約され、他のシーケンス長やバッチサイズでは動作しません。異なるサイズで試すと、以下のエラーが発生します：

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

お勧めしますのは、モデルの推論中に供給される最大の入力と同じ大きさのダミー入力サイズでモデルをトレースすることです。パディングを使用して不足値を補完することもできます。ただし、モデルがより大きな入力サイズでトレースされるため、行列の寸法も大きくなり、より多くの計算が発生します。

異なるシーケンス長のモデルをエクスポートする際に、各入力に対して実行される演算の総数に注意して、パフォーマンスを密接にフォローすることをお勧めします。

##
Using TorchScript in Python

このセクションでは、モデルの保存と読み込み、および推論にトレースを使用する方法を示します。

###
Saving a model

TorchScriptで`BertModel`をエクスポートするには、`BertConfig`クラスから`BertModel`をインスタンス化し、それをファイル名`traced_bert.pt`でディスクに保存します：

    from transformers import BertModel, BertTokenizer, BertConfig
    import torch

    enc = BertTokenizer.from_pretrained("google-bert/bert-base-uncased")

    # Tokenizing input text
    text = "[CLS] Who was Jim Henson ? [SEP] Jim Henson was a puppeteer [SEP]"
    tokenized_text = enc.tokenize(text)

    # Masking one of the input tokens
    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]

    # Creating a dummy input
    tokens_tensor = torch.tensor([indexed_tokens])
    segments_tensors = torch.tensor([segments_ids])
    dummy_input = [tokens_tensor, segments_tensors]

    # Initializing the model with the torchscript flag
    # Flag set to True even though it is not necessary as this model does not have an 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,
    )

    # Instantiating the model
    model = BertModel(config)

    # The model needs to be in evaluation mode
    model.eval()

    # If you are instantiating the model with *from_pretrained* you can also easily set the TorchScript flag
    model = BertModel.from_pretrained("google-bert/bert-base-uncased", torchscript=True)

    # Creating the trace
    traced_model = torch.jit.trace(model, [tokens_tensor, segments_tensors])
    torch.jit.save(traced_model, "traced_bert.pt")

###
Loading a model

以前に保存した `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)

###
Using a traced model for inference

トレースモデルを使用して推論を行うには、その `__call__` ダンダーメソッドを使用します。

    traced_model(tokens_tensor, segments_tensors)

##
Deploy Hugging Face TorchScript models to AWS with the Neuron SDK

AWSはクラウドでの低コストで高性能な機械学習推論向けに Amazon EC2 Inf1 インスタンスファミリーを導入しました。Inf1インスタンスはAWS Inferentiaチップによって駆動され、ディープラーニング推論ワークロードに特化したカスタムビルドのハードウェアアクセラレータです。AWS Neuron はInferentia用のSDKで、トランスフォーマーモデルをトレースして最適化し、Inf1に展開するためのサポートを提供します。

Neuron SDK が提供するもの:

1. クラウドでの推論のためにTorchScriptモデルをトレースして最適化するための、1行のコード変更で使用できる簡単なAPI。
2. 改善されたコストパフォーマンス のためのボックス外のパフォーマンス最適化。
3. PyTorch または TensorFlow で構築されたHugging Faceトランスフォーマーモデルへのサポート。

###
Implications

BERT（Bidirectional Encoder Representations from Transformers）アーキテクチャやその変種（distilBERT や roBERTa など）に基づくトランスフォーマーモデルは、非生成タスク（抽出型質問応答、シーケンス分類、トークン分類など）において、Inf1上で最適に動作します。ただし、テキスト生成タスクも AWS Neuron MarianMT チュートリアル に従ってInf1上で実行できます。Inferentiaでボックス外で変換できるモデルに関する詳細情報は、Neuronドキュメンテーションの Model Architecture Fit セクションにあります。

###
Dependencies

モデルをAWS Neuronに変換するには、Neuron SDK 環境 が必要で、AWS Deep Learning AMI に事前に構成されています。

###
Converting a model for AWS Neuron

モデルをAWS NEURON用に変換するには、PythonでTorchScriptを使用する と同じコードを使用して `BertModel` をトレースします。Python APIを介してNeuron SDKのコンポーネントにアクセスするために、`torch.neuron` フレームワーク拡張をインポートします。

    from transformers import BertModel, BertTokenizer, BertConfig
    import torch
    import torch.neuron

次の行を変更するだけで済みます。

    - torch.jit.trace(model, [tokens_tensor, segments_tensors])
    + torch.neuron.trace(model, [token_tensor, segments_tensors])

これにより、Neuron SDKはモデルをトレースし、Inf1インスタンス向けに最適化します。

AWS Neuron SDKの機能、ツール、サンプルチュートリアル、最新のアップデートについて詳しく知りたい場合は、AWS NeuronSDK ドキュメンテーション をご覧ください。

< > Update on GitHub