AI 物見遊山 - 自然言語処理編 - LSTM モデル推論の効率化 その壱

AI 物見遊山トップページへ

自然言語処理編 - LSTM モデル推論の効率化 その壱

アイキャッチ画像

前の記事では、推論の正解率改善を行った。 この記事では、推論の効率化を体験する。 効率化は、モデルの ONNX 形式変換で行う。 ONNX (Open Neural Network Exchange) とは、 機械学習モデルの形式を定義する標準規格であり、 異なる環境間でモデルを共有したり、実行したりすることが容易になる。 ONNX 形式のモデルは、 ONNX Runtime 上で高速に実行できる。 ONNX Runtime は、様々なプラットフォーム (Windows, Linux, macOS, iOS, Android など) や、 ハードウェア (CPU, GPU, FPGA など) 上で、 ONNX 形式のモデルを実行するための推論エンジンとなる。

ONNX Runtime をインストールする

各パッケージは、複数インストールして使い分けることができる。

CPU で動作させる場合

pip install onnxruntime onnx

CUDA で動作させる場合

pip install onnxruntime-gpu onnx

DirectML で動作させる場合

DirectML (Direct Machine Learning) は、 Microsoft 製の機械学習向けの低水準 API となる。 DirectX 12 スタイルで、機械学習の推論を行える。 モデル学習はできない。 AMD 製の GPU でも利用できる。 現時点では、このパッケージを WSL 上にインストールできないので注意。

pip install onnxruntime-directml onnx

LSTM モデルを ONNX 形式に変換する

opset_version は、 ONNX Versioning に記載がある。 ただ、新しい程よいわけでなく、推論環境でサポートしている必要がある。 ハードウェアや、ドライバのバージョンによっては、期待した動作にならないようだ。

convert2onnx.py

import bz2
import logging
import os

import onnx
import time
import torch

from pos_tagger_lstm import PosTaggerLSTM

# ログ設定
logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s', level=logging.INFO)

# データディレクトリ
DATA_DIR = '.'

# 処理済みの単語を重複なく保持する入出力ファイル名
WORD_FILE = os.path.join(DATA_DIR, 'lstm-word.txt.bz2')
# 処理済みの品詞を重複なく保持する入出力ファイル名
POS_FILE = os.path.join(DATA_DIR, 'lstm-pos.txt.bz2')

# テスト用訓練データ出力ファイル名
X_TEST_FILE = os.path.join(DATA_DIR, 'lstm-optim-x-test.npy')

# モデルファイル名
MODEL_FILE = os.path.join(DATA_DIR, 'lstm-layer2-bidirect-epoch10.model')
# 変換後モデルファイル名
ONNX_FILE = os.path.join(DATA_DIR, 'lstm-layer2-bidirect.onnx')


def main():
    """LSTM モデル学習を行う"""
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    print('Loading datas...')
    # 単語 ID 辞書を読み込む
    with bz2.open(WORD_FILE, 'rt', encoding='utf-8') as f:
        word_ids = {k: int(v) for k, v in (line.split('\t') for line in f)}
    # 品詞 ID 辞書を読み込む
    with bz2.open(POS_FILE, 'rt', encoding='utf-8') as f:
        pos_ids = {k: int(v) for k, v in (line.split('\t') for line in f)}

    print('Loading model...')
    # PyTorch モデルを準備する
    model = PosTaggerLSTM(
        len(word_ids) + 1,
        len(pos_ids),
        100,
        num_layers=2,
        bidirectional=True
    ).to(device)
    checkpoint = torch.load(MODEL_FILE, weights_only=True, map_location=device)
    model.load_state_dict(checkpoint['model_state_dict'])

    print('Converting to ONNX...')
    # PyTorch モデルを ONNX モデルに変換する
    torch.onnx.export(
        model,
        # ダミー入力。引数は、先頭から 0 から 1 の範囲で、(バッチサイズ、データサイズ) という指定となる。
        # バッチサイズと、データサイズは、動的サイズとするため、ここではそれぞれ 1 にしている。
        # 動的サイズ指定をしない場合、推論時に同じサイズのデータ以外は、エラーとなる。
        torch.randint(0, 1, (1, 1)).to(device),
        ONNX_FILE,
        opset_version=20,
        # 動的パラメータ(dynamic_axes)を指定するために設定が必要となる。
        input_names=['input'],
        # output_names は、複数の出力があり、推論時に選択して取得したい場合などに設定する。
        output_names=['output'],
        dynamic_axes={
            # バッチサイズとデータサイズを可変にする
            # インデックスと名称を指定する。名称は、任意の名前とする.
            'input': {0: 'batch_size', 1: 'data_dim'},
            'output': {0: 'batch_size', 1: 'data_dim'}
        }
    )
    print(f'ONNX model size: {os.path.getsize(ONNX_FILE) / 1024 / 1024:.2f} MB')

    print('Checking ONNX model...')
    onnx_model = onnx.load(ONNX_FILE)
    onnx.checker.check_model(onnx_model)
    print('ONNX model is valid.')


if __name__ == '__main__':
    start_time = time.time()
    main()
    print(f'経過時間: {time.time() - start_time:.2f} s')

以下のような実行結果となった。 CPU で実行している。

Loading datas...
Loading model...
Converting to ONNX...
/home/ubuntu/miniconda3/envs/ai/lib/python3.12/site-packages/onnxscript/converter.py:823: FutureWarning:

'onnxscript.values.Op.param_schemas' is deprecated in version 0.1 and will be removed in the future. Please use '.op_signature' instead.

/home/ubuntu/miniconda3/envs/ai/lib/python3.12/site-packages/onnxscript/converter.py:823: FutureWarning:

'onnxscript.values.OnnxFunction.param_schemas' is deprecated in version 0.1 and will be removed in the future. Please use '.op_signature' instead.

/home/ubuntu/miniconda3/envs/ai/lib/python3.12/site-packages/torch/onnx/symbolic_opset9.py:4545: UserWarning:

Exporting a model to ONNX with a batch_size other than 1, with a variable length with LSTM can cause an error when running the ONNX model with a different batch size. Make sure to save the model with a batch size of 1, or define the initial states (h0/c0) as inputs of the model. 

ONNX model size: 17.56 MB
Checking ONNX model...
ONNX model is valid.
経過時間: 1.18 s

ONNX モデルで推論をする

ループ回数、初回ループ処理時間や、次回ループ以降の平均処理時間などを表示するようにした。 初回ループ処理は、初期処理等で遅いらしいので、個別に確認できるようにした。

infer_onnx.py

import logging
import os

import numpy as np
import onnxruntime
import time

# ログ設定
logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s', level=logging.INFO)

# データディレクトリ
DATA_DIR = '.'

# 処理済みの単語を重複なく保持する入出力ファイル名
WORD_FILE = os.path.join(DATA_DIR, 'lstm-word.txt.bz2')
# 処理済みの品詞を重複なく保持する入出力ファイル名
POS_FILE = os.path.join(DATA_DIR, 'lstm-pos.txt.bz2')

# テスト用訓練データ出力ファイル名
X_TEST_FILE = os.path.join(DATA_DIR, 'lstm-optim-x-test.npy')
# テスト用教師データ出力ファイル名
Y_TEST_FILE = os.path.join(DATA_DIR, 'lstm-optim-y-test.npy')

# モデルファイル名
MODEL_FILE = os.path.join(DATA_DIR, 'lstm-layer2-bidirect.onnx')

# バッチサイズ
BATCH_SIZE = 512


def main() -> None:
    """
    LSTM ONNX モデル推論を行う
    """
    available_providers = onnxruntime.get_available_providers()
    providers = ['CUDAExecutionProvider'] if 'CUDAExecutionProvider' in available_providers else [
        'CPUExecutionProvider']
    print(f'Using providers: {providers}')

    print('Loading datas...')
    # テスト用訓練データファイルを読み込む
    # x_test = np.load(X_TEST_FILE, mmap_mode='r')
    x_test = np.load(X_TEST_FILE)
    # テスト用教師データファイルを読み込む
    # y_test = np.load(Y_TEST_FILE, mmap_mode='r')
    y_test = np.load(Y_TEST_FILE)

    options = onnxruntime.SessionOptions()
    options.enable_profiling = False
    ort_session = onnxruntime.InferenceSession(MODEL_FILE, options, providers=providers)
    loop_count = 0
    total = 0
    first_time = 0
    ok = 0
    print('Start inference...')
    start_time = time.time()
    for i in range(0, len(x_test), BATCH_SIZE):
        loop_count += 1
        x_batch = x_test[i:i + BATCH_SIZE]
        y_batch = y_test[i:i + BATCH_SIZE]

        # パディングなら False、それ以外は True となる配列を作成する
        mask = y_batch != -1
        # パディング以外のデータ数を集計する(True のみ集計対象となる)
        total += mask.sum()
        # 推論をする
        ort_outputs = ort_session.run(
            None,
            {'input': x_batch}
        )
        output = ort_outputs[0]

        ans = np.argmax(output, axis=2)
        # パディングデータ以外の正答数を集計する
        ok += ((ans == y_batch) * mask).sum()
        if i == 0:
            first_time = time.time()

    profile_file = ort_session.end_profiling()
    print(f'Profile file: {profile_file}')

    print(f'OK: {ok}, Total: {total}, Accuracy: {ok / total:.4f}')
    elapsed_time = time.time() - start_time
    print(f'全体経過時間: {elapsed_time:.4f} s ({total / elapsed_time:.0f} times/s)')
    print(f'ループ回数: {loop_count}')
    first_elapsed_time = first_time - start_time
    print(f'初回時間: {first_elapsed_time:.4f} s')
    print(f'次回以降平均時間: {(elapsed_time - first_elapsed_time) / (loop_count - 1):.4f} s')


if __name__ == '__main__':
    main()

以下のような実行結果となった。

CPU で実行した場合

Using providers: ['CPUExecutionProvider']
Loading datas...
Start inference...
Profile file: 
OK: 428282, Total: 440767, Accuracy: 0.9717
全体経過時間: 7.5133 s (58665 times/s)
ループ回数: 27
初回時間: 0.4589 s
次回以降平均時間: 0.2713 s

T4 GPU で実行した場合

Using providers: ['CUDAExecutionProvider']
Loading datas...
Start inference...
Profile file: 
OK: 428196, Total: 440767, Accuracy: 0.9715
全体経過時間: 3.2806 s (134355 times/s)
ループ回数: 27
初回時間: 0.2305 s
次回以降平均時間: 0.1173 s

まとめ

今回の処理速度の比較は、以下のようになった。 ただし、簡易的なモデルであることや、パラメータなども未調整なので、あくまで参考値となる。

モデル CPU(秒) T4 GPU(秒)
PyTorch 16.66 4.96
ONNX 7.51 (2.22 倍速) 3.28 (1.51 倍速)

PyTorch(CPU) -> ONNX(T4 GPU) で、5.1 倍速程度。

単純な変換のみで、結構な効率化となった。 次の記事では、さらに量子化を行い、追加の効率化をしてみる。

参考資料

場所: 日本、東京都武蔵野市