PyTorch RNN-T モデルの変換#
危険
ここで説明されているコードは非推奨になりました。従来のソリューションの適用を避けるため使用しないでください。下位互換性を確保するためにしばらく保持されますが、最新のアプリケーションでは使用してはなりません。
このガイドでは、非推奨となった変換方法について説明します。新しい推奨方法に関するガイドは、Python チュートリアルに記載されています。
このガイドでは、MLCommons リポジトリーからの RNN-T モデルの変換について説明します。IR に変換する前に、以下の手順に従って PyTorch モデルを ONNX にエクスポートします:
ステップ 1.MLCommons リポジトリー (リビジョン r1.0) から RNN-T PyTorch 実装のクローンを作成します。完全なリポジトリーを使用せずに RNN-T モデルのみを取得する浅いクローンを作成します。すでに完全なリポジトリーがある場合、これをスキップして ステップ 2 に進みます:
git clone -b r1.0 -n https://github.com/mlcommons/inference rnnt_for_openvino --depth 1
cd rnnt_for_openvino
git checkout HEAD speech_recognition/rnntステップ 2.MLCommons 推論リポジトリーの完全なクローンがすでに存在する場合、IR への変換が行われる事前トレーニング済みの PyTorch モデル用のフォルダーを作成します。ステップ 5 では、完全なクローンへのパスも指定する必要があります。浅いクローンがある場合は、この手順をスキップしてください。
mkdir rnnt_for_openvino
cd rnnt_for_openvinoステップ 3.こちらから、PyTorch 実装用の事前トレーニングされた重みをダウンロードします。UNIX のようなシステムでは wget を使用できます:
wget https://zenodo.org/record/3662521/files/DistributedDataParallel_1576581068.9962234-epoch-100.ptリンクは speech_recoginitin/rnnt サブフォルダー内の setup.sh から取得されました。ガイドに従った場合と全く同じ重みが得られます。
ステップ 4.必要な Python パッケージをインストールします:
pip3 install torch tomlステップ 5.以下のスクリプトを使用して、RNN-T モデルを ONNX にエクスポートします。以下のコードを export_rnnt_to_onnx.py ファイルにコピーし、現在のディレクトリー rnnt_for_openvino で実行します。
注
MLCommons 推論リポジトリーの完全なクローンがすでにある場合は、mlcommons_inference_path 変数を指定する必要があります。
import toml
import torch
import sys
def load_and_migrate_checkpoint(ckpt_path):
checkpoint = torch.load(ckpt_path, map_location="cpu") }
migrated_state_dict = {}
for key, value in checkpoint['state_dict'].items():
key = key.replace("joint_net", "joint.net")
migrated_state_dict[key] = value
del migrated_state_dict["audio_preprocessor.featurizer.fb"]
del migrated_state_dict["audio_preprocessor.featurizer.window"]
return migrated_state_dict
mlcommons_inference_path = './'# specify relative path for MLCommons inferene
checkpoint_path = 'DistributedDataParallel_1576581068.9962234-epoch-100.pt'
config_toml = 'speech_recognition/rnnt/pytorch/configs/rnnt.toml'
config = toml.load(config_toml)
rnnt_vocab = config['labels']['labels']
sys.path.insert(0, mlcommons_inference_path + 'speech_recognition/rnnt/pytorch')
from model_separable_rnnt import RNNT
model = RNNT(config['rnnt'], len(rnnt_vocab) + 1, feature_config=config['input_eval'])
model.load_state_dict(load_and_migrate_checkpoint(checkpoint_path))
seq_length, batch_size, feature_length = 157, 1, 240
inp = torch.randn([seq_length, batch_size, feature_length])
feature_length = torch.LongTensor([seq_length])
x_padded, x_lens = model.encoder(inp, feature_length)
torch.onnx.export(model.encoder, (inp, feature_length), "rnnt_encoder.onnx", opset_version=12,
input_names=['input', 'feature_length'], output_names=['x_padded', 'x_lens'],
dynamic_axes={'input': {0: 'seq_len', 1: 'batch'}})
symbol = torch.LongTensor([[20]])
hidden = torch.randn([2, batch_size, 320]), torch.randn([2, batch_size, 320])
g, hidden = model.prediction.forward(symbol, hidden)
torch.onnx.export(model.prediction, (symbol, hidden), "rnnt_prediction.onnx", opset_version=12,
input_names=['symbol', 'hidden_in_1', 'hidden_in_2'], output_names=['g', 'hidden_out_1', 'hidden_out_2'],
dynamic_axes={'symbol': {0: 'batch'}, 'hidden_in_1': {1: 'batch'}, 'hidden_in_2': {1: 'batch'}})
f = torch.randn([batch_size, 1, 1024])
model.joint.forward(f, g)
torch.onnx.export(model.joint, (f, g), "rnnt_joint.onnx", opset_version=12,
input_names=['0', '1'], output_names=['result'], dynamic_axes={'0': {0: 'batch'}, '1': {0: 'batch'}})python3 export_rnnt_to_onnx.pyこの手順が完了すると、ファイル rnnt_encoder.onnx、rnnt_prediction.onnx、および rnnt_joint.onnx が現在のディレクトリーに保存されます。
ステップ 6.変換コマンドを実行します:
mo --input_model rnnt_encoder.onnx --input "input[157,1,240],feature_length->157"
mo --input_model rnnt_prediction.onnx --input "symbol[1,1],hidden_in_1[2,1,320],hidden_in_2[2,1,320]"
mo --input_model rnnt_joint.onnx --input "0[1,1,1024],1[1,1,320]"注
シーケンス長 = 157 のハードコードされた値は MLCommons から取得されたものですが、IR への変換によりネットワークの再形成の可能性が維持されます。入力形状は、変換中または推論中に手動で任意の値に変更できます。