1. torchscript と jit の概要
1. トーチスクリプトについて
TorchScript は、Pytorch モデルの中間表現です (nn.Module から継承)。保存された torchscript モデルは、C++ のような高性能環境で実行できます。
TorchScript は、PyTorch コードからシリアル化可能で最適化可能なモデルを作成する方法です。どのようなトーチランプ プログラムも Python プロセスから保存し、Python 依存関係のないプロセスにロードできます。
簡単に言うと、TorchScript は動的グラフを静的グラフに変換することができ、pytorch の柔軟な動的グラフ機能を利用してモデル構造 (モデル定義) を取得できるツールを提供します。
2.torch.jitについて
JITとは何ですか?
まず第一に、JIT は概念であることを知っておく必要があります。正式名は Just In Time Compilation で、中国語では「ジャスト イン タイム コンパイル」と訳されます。プログラムの最適化手法です。一般的な使用シナリオは「通常のコンパイル」です。表現"。たとえば、Python で正規表現を使用すると、次のようになります。
prog = re.compile(pattern)
result = prog.match(string)
#或
result = re.match(pattern, string)
上記 2 つの例は Python の公式ドキュメントから直接抽出したもので、ドキュメントからは 2 つの記述方法が結果の点で「同等」であることがわかります。ただし、最初の書き方に注意してください。正規表現は最初にコンパイルされてから使用されます。Python ドキュメントを読み続けると、次の文章を見つけることができます。
re.compile() を使用し、結果の正規表現オブジェクトを再利用のために保存すると、正規表現が 1 つのプログラムで複数回使用される場合に効率的になります。そして、このコンパイル処理は、JIT (Just-In-Time Compilation) として理解できます。
PyTorch は発売以来その「使いやすさ」で知られており、PyTorch の「動的グラフ」構造によりネイティブ Python の開発に最適です。PyTorch モデルの前に任意の Python プロセス制御文を追加することができ、ブレークポイントを経由しても問題ありませんが、TensorFlow の場合は tf.cond コントロールなど TensorFlow で開発されたプロセスを使用する必要があります。動的グラフ モデルは、使いやすさのためにいくつかの高度な機能を犠牲にします。
JIT の利点:
1. モデルのデプロイメント
PyTorch バージョン 1.0 でリリースされた 2 つのコア新機能は、JIT と C++ API です。これら 2 つの機能を一緒にリリースすることは不合理ではありません。JIT は Python と C++ の間のブリッジです。Python を使用してモデルをトレーニングできます。その後、モデルは JIT によって言語に依存しないモジュールに変換されるため、C++ を非常に簡単に呼び出すことができるようになり、これ以降、「Python を使用してモデルをトレーニングし、C++ を使用してモデルを本番環境にデプロイする」ことが簡単になりました。 PyTorch のタスク。また、C++ を使用しているため、Raspberry Pi、iOS、Android など、ほぼすべてのプラットフォームやデバイスに PyTorch モデルをデプロイできるようになりました。
- パフォーマンスの向上
デプロイおよび運用のために提供される機能であるため、パフォーマンスに関して大幅な最適化が行われていることは避けられません。推論されたシーンに高いパフォーマンス要件がある場合は、モデル (torch.nn.Module) を TorchScript モジュールに変換することを検討できます。それから推論を続けます。
- モデルの視覚化
TensorFlow または Keras は、静的グラフ プログラミング モデルであるため、モデル視覚化ツール (TensorBoard など) に非常に適しています。モデルが定義された後、モデル全体の構造とフォワード ロジックはすでに明確になっていますが、PyTorch 自体は明確ではありません。そのため、PyTorch モデルは視覚化が常に苦手でしたが、JIT によって状況が改善されました。これで、JIT のトレース機能を使用して、特定の入力に対する PyTorch モデルのフォワード ロジックを取得できるようになり、フォワード ロジックを通じてモデルのおおよその構造を取得できます。(ただし、forwardメソッド内に多くの条件付き制御ステートメントがある場合、これはまだ良い方法ではありません)
3. トーチスクリプトモジュールを生成する 2 つの方法
1. スクリプト作成
TorchScript 言語を直接使用して PyTorch JIT モジュールを定義し、次に torch.jit.script を使用してそれを TorchScript モジュールに変換し、ファイルとして保存できます。TorchScript 言語自体も Python コードであるため、Python ファイルに直接記述することができます。
TorchScript 言語の使用は TensorFlow の使用と似ており、事前に完全なグラフを定義する必要があります。TensorFlow の場合、条件制御に Python の if ステートメントやその他のステートメントを直接使用できないことがわかっており、tf.cond を使用する必要があります。しかし、TorchScript の場合は、静的グラフでも if や for などの条件制御ステートメントを直接使用できます。 , PyTorchは今でも「使いやすさ」にこだわり続けています。TorchScript 言語は、Python の静的に型指定されたサブセットであり、Python 3 の型指定モジュールを使用して実装されるため、TorchScript 言語の作成エクスペリエンスは、一部の Python 機能が使用できないことを除いて、Python のエクスペリエンスとまったく同じです (これを「TorchScript 言語リファレンス」に渡すと、ネイティブ Python との類似点と相違点を確認できます。
理論的には、グラフ構造全体が事前に定義されているため、スクリプティングによって定義された TorchScript モジュールはモデル視覚化ツールに非常に適しています。
- トレース
TorchScript モジュールを使用するより簡単な方法は、PyTorch モデル (torch.nn.Module) を TorchScript モジュールに直接変換できる Tracing を使用することです。「トラッキング」とは、その名のとおり、モデルを再度前進させるための「入力」を提供し、入力のフローパスを通じてグラフの構造を取得することです。この方法は、単純な前方ロジックを持つモデルでは非常に実用的ですが、前方自体に多くのフロー制御ステートメントがある場合、同じ入力がすべての論理分岐を横断できないため、問題が発生する可能性があります。
2. 推論用のトーチ モデルを生成する
1. エクスポートされたトーチ チェックポイント モデルをロードします
事前トレーニングされたモデル構成ファイルと書き換えられたモデル構造をロードします
# 【multitask_classify_ner 多任务分类模型代码（包括classify任务和ner任务）】
class BertFourLevelArea(BertPreTrainedModel):
"""BERT model for four level area.
"""
def __init__(self, config, num_labels_cls, num_labels_ner, inner_dim, RoPE):
super(BertFourLevelArea, self).__init__(config, num_labels_cls, num_labels_ner, inner_dim, RoPE)
self.bert = BertModel(config)
self.num_labels_cls = num_labels_cls
self.num_labels_ner = num_labels_ner
self.inner_dim = inner_dim
self.hidden_size = config.hidden_size
self.dense_ner = nn.Linear(self.hidden_size, self.num_labels_ner * self.inner_dim * 2)
self.dense_cls = nn.Linear(self.hidden_size, num_labels_cls)
self.RoPE = RoPE
self.dropout = nn.Dropout(config.hidden_dropout_prob)
self.apply(self.init_bert_weights)
def sinusoidal_position_embedding(self, batch_size, seq_len, output_dim):
position_ids = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(-1)
indices = torch.arange(0, output_dim // 2, dtype=torch.float)
indices = torch.pow(10000, -2 * indices / output_dim)
embeddings = position_ids * indices
embeddings = torch.stack([torch.sin(embeddings), torch.cos(embeddings)], dim=-1)
embeddings = embeddings.repeat((batch_size, *([1]*len(embeddings.shape))))
embeddings = torch.reshape(embeddings, (batch_size, seq_len, output_dim))
embeddings = embeddings.to(self.device)
return embeddings
def forward(self, input_ids, token_type_ids=None, attention_mask=None):
# sequence_output: Last Encoder Layer.shape: (batch_size, seq_len, hidden_size)
encoded_layers, pooled_output = self.bert(input_ids, token_type_ids, attention_mask)
sequence_output = encoded_layers[-1]
batch_size = sequence_output.size()[0]
seq_len = sequence_output.size()[1]
# 【Bert Ner GlobalPointer】:
# outputs: (batch_size, seq_len, num_labels_ner*inner_dim*2)
outputs = self.dense_ner(sequence_output)
# outputs: (batch_size, seq_len, num_labels_ner, inner_dim*2)
outputs = torch.split(outputs, self.inner_dim * 2, dim=-1) # TODO:1
outputs = torch.stack(outputs, dim=-2) # TODO:2
# qw,kw: (batch_size, seq_len, num_labels_ner, inner_dim)
qw, kw = outputs[...,:self.inner_dim], outputs[...,self.inner_dim:] # TODO:3
if self.RoPE:
# pos_emb:(batch_size, seq_len, inner_dim)
pos_emb = self.sinusoidal_position_embedding(batch_size, seq_len, self.inner_dim)
# cos_pos,sin_pos: (batch_size, seq_len, 1, inner_dim)
cos_pos = pos_emb[..., None, 1::2].repeat_interleave(2, dim=-1)
sin_pos = pos_emb[..., None,::2].repeat_interleave(2, dim=-1)
qw2 = torch.stack([-qw[..., 1::2], qw[...,::2]], -1)
qw2 = qw2.reshape(qw.shape)
qw = qw * cos_pos + qw2 * sin_pos
kw2 = torch.stack([-kw[..., 1::2], kw[...,::2]], -1)
kw2 = kw2.reshape(kw.shape)
kw = kw * cos_pos + kw2 * sin_pos
# logits_ner:(batch_size, num_labels_ner, seq_len, seq_len)
logits_ner = torch.einsum('bmhd,bnhd->bhmn', qw, kw) # TODO:4
# padding mask
pad_mask = attention_mask.unsqueeze(1).unsqueeze(1).expand(batch_size, self.num_labels_ner, seq_len, seq_len) # TODO:5
# pad_mask_h = attention_mask.unsqueeze(1).unsqueeze(-1).expand(batch_size, self.num_labels_ner, seq_len, seq_len)
# pad_mask = pad_mask_v&pad_mask_h
logits_ner = logits_ner*pad_mask - (1-pad_mask)*1e12 # TODO:6
# 排除下三角
mask = torch.tril(torch.ones_like(logits_ner), -1) # TODO:7
logits_ner = logits_ner - mask * 1e12 # TODO:8
# 【Bert Classify】:
pooled_output = self.dropout(pooled_output)
logits_cls = self.dense_cls(pooled_output)
return logits_cls, logits_ner
#【加载预训练模型参数】
config = modeling.BertConfig.from_json_file('/root/ljh/space-based/Deep_Learning/Pytorch/multitask_classify_ner/pretrain_model/bert-base-chinese/config.json')
#【加载我们训练的模型 】
#【num_labels_cls 和 num_labels_ner为我们训练的label_counts 这次训练的分类任务标签数为1524 ，NER任务的分类数为13】
num_labels_cls = 1524
num_labels_ner = 13
model = modeling.BertFourLevelArea(
config,
num_labels_cls=num_labels_cls,
num_labels_ner=num_labels_ner,
inner_dim=64,
RoPE=False
)
2. モデルモデルパラメータのロード
トレーニング済みモデルのパラメータをロードする
#【训练完成的tourch 模型地址】
init_checkpoint='/root/ljh/space-based/Deep_Learning/Pytorch/multitask_classify_ner/outputs.bak/multitask_classify_ner/pytorch_model.bin'
#【载入模型】
checkpoint = torch.load(init_checkpoint, map_location=torch.device("cuda"))
checkpoint = checkpoint["model"] if "model" in checkpoint.keys() else checkpoint
model.load_state_dict(checkpoint)
device = torch.device("cuda")
#【将模型导入GPU】
model = model.to(device)
#【模型初始化】
model.eval()
3. torch.jit 追跡パラメータを構築する
アドレスクリーニングマルチタスクモデルの順伝播ロジックには複数の判定条件構造がないため、順伝播のプロセスとモデルの構造を記録するためにトレース（Tracing）の形式を選択します。
#【定义tokenizer 】
from transformers import BertTokenizerFast
tokenizer = BertTokenizerFast.from_pretrained('/root/ljh/space-based/Deep_Learning/Pytorch/multitask_classify_ner/pretrain_model/bert-base-chinese', add_special_tokens=True, do_lower_case=False)
input_str='上海上海市青浦区华隆路E通世界华新园'
max_seq_length=64
#【生成bert模型输出】
def input2feature(input_str, max_seq_length=48):
# 预处理字符
tokens_a = tokenizer.tokenize(input_str)
# 如果超过长度限制，则进行截断
if len(input_str) > max_seq_length - 2:
tokens_a = tokens_a[0:(max_seq_length - 2)]
tokens = ["[CLS]"] + tokens_a + ["[SEP]"]
input_ids = tokenizer.convert_tokens_to_ids(tokens)
input_length = len(input_ids)
input_mask = [1] * input_length
segment_ids = [0] * input_length
while len(input_ids) < max_seq_length:
input_ids.append(0)
input_mask.append(0)
segment_ids.append(0)
return input_ids, input_mask, segment_ids
#【输入地址token化 input_ids --> list()】
input_ids, input_mask, segment_ids = input2feature(input_str, max_seq_length)
# 【list -> tensor】
input_ids = torch.tensor(input_ids, dtype=torch.long)
input_mask = torch.tensor(input_mask, dtype=torch.long)
segment_ids = torch.tensor(segment_ids, dtype=torch.long)
#【这里stack 是因为模型内部定义的输出参数需要stack 】
input_ids = torch.stack([input_ids], dim=0)
input_mask = torch.stack([input_mask], dim=0)
segment_ids = torch.stack([segment_ids], dim=0)
#【将参数推送至cuda设备中】
device = torch.device("cuda")
input_ids = input_ids.to(device)
input_mask = input_mask.to(device)
segment_ids = segment_ids.to(device)
#【input_ids.shape --> torch.Size([1, 64])】
4. torch.jit を使用して、TorchScript モジュール モデルをエクスポートします。
jit はトレース (Tracing) の形式を使用して、順伝播プロセスとモデルの構造を記録します。
#【根据输出的input_ids, input_mask, segment_id记录前向传播过程】
script_model = torch.jit.trace(model,[input_ids, input_mask, segment_ids],strict=True)
#【保存】
torch.jit.save(script_model, "./multitask_test/multitask_model/1/model.pt")
5.TorchScript モジュールが正しいことを確認します
#【查看torch模型结果】
cls_res, ner_res = model(input_ids, input_mask, segment_ids)
import numpy as np
np.argmax(cls_res.detach().cpu().numpy())
#【result:673】
#【load torchscript model】
jit_model = torch.jit.load('./multitask_test/multitask_model/1/model.pt')
example_outputs_cls,example_outputs_ner = jit_model(input_ids, input_mask, segment_ids)
np.argmax(example_outputs_cls.detach().cpu().numpy())
#【result:673】
3. triton サーバーを使用して torchscript モデルを起動します
1. config.ptxtx 構成ファイルを変更します。
name: "multitask_model"
platform: "pytorch_libtorch"
max_batch_size: 8
input [
{
name: "input_ids"
data_type: TYPE_INT64
dims: 64
},
{
name: "segment_ids"
data_type: TYPE_INT64
dims: 64
},
{
name: "input_mask"
data_type: TYPE_INT64
dims: 64
}
]
output [
{
name: "cls_logits"
data_type: TYPE_FP32
dims: [1, 1524]
},
{
name: "ner_logits"
data_type: TYPE_FP32
dims: [ -1, 13, 64, 64 ]
}
]
dynamic_batching {
preferred_batch_size: [ 1, 2, 4, 8 ]
max_queue_delay_microseconds: 50
}
instance_group [
{
count: 1
kind: KIND_GPU
gpus: [0]
}
]
2. モデルのディレクトリ構造
multitask_test
∴── multitask_model
§── 1
│ └── model.pt
§── config.pbtxt
└── label_cls.csv
3.トリトンを起動する
tritonserver --model-store=/root/ljh/space-based/Deep_Learning/Pytorch/multitask_classify_ner/multitask_test --strict-model-config=false --exit-on-error=false
4. 解決すべきTorchscriptモデルモデルのデプロイメント
torchscript マルチカード デプロイメントを使用しようとするプロセスで、モデルモデルが cuda にバインドされます。triton でモデルを起動した後、モデルにバインドされた cuda デバイスのみが正常に実行できます。
同様の問題のリファレンス: https://github.com/triton-inference-server/server/issues/2626