自然言語処理編 - LSTM モデル学習その参

前の記事では、モデル学習を行った。ただし、非効率な実装となっていたため、この記事では、モデル学習の効率化をしてみる。モデル学習の効率化として以下の内容を行う。

バッチ処理化：訓練データを一定量まとめて学習する。
学習率を動的に調整する。

バッチ処理化には、学習データの配列サイズを統一する必要がある。一文に含まれる単語数は文によって異なるため、以下のように配列長が異なっている。

[
  [  5,  3],
  [ 10, 99, 15,  7],
  [ 20, 22,  6]
]

上記のようなデータを、パディング処理をして以下のように変換する必要がある。

[
  [  5,  3,  0,  0],
  [ 10, 99, 15,  7],
  [ 20, 22,  6,  0]
]

パディング値は、訓練データの場合 0、教師データの場合は、-1 とする。

学習率の動的変更には、スケジューラを導入する。スケジューラには、いくつかの種類があり、状況によって使い分ける必要がありそうだ。

学習用データ生成処理を変更する

以前の記事で作成した学習用データ生成処理を変更する。要点は、データをパディング処理し、Numpy 形式でファイル保存するようにする。生成された学習用データは、パディング処理を行ったため、ファイルサイズがかなり増えている。

generate_train_data_optim.py

import bz2
import os
from collections import defaultdict

import MeCab
import numpy as np
import torch
from torch.nn.utils.rnn import pad_sequence

# 学習用文章データの入力ファイル名
TRAIN_FILE = 'jawiki-data.txt.bz2'
# テスト用文章データの入力ファイル名
TEST_FILE = 'jawiki-test.txt.bz2'

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

# 学習用訓練データ出力ファイル名
X_TRAIN_FILE = 'lstm-optim-x-train.npy'
# 学習用教師データ出力ファイル名
Y_TRAIN_FILE = 'lstm-optim-y-train.npy'
# テスト用訓練データ出力ファイル名
X_TEST_FILE = 'lstm-optim-x-test.npy'
# テスト用教師データ出力ファイル名
Y_TEST_FILE = 'lstm-optim-y-test.npy'


def load_ids(file_path: str) -> dict[str, int]:
    """ID 辞書を読み込む。

    Args:
        file_path (str): ファイルパス

    Returns:
        dict[str, int]: ID 辞書
    """
    # インデックス 0 は、データのパディング用に利用するため、空文字で初期化しておく。
    ids = {'': 0}
    if os.path.isfile(file_path):
        with bz2.open(file_path, 'rt', encoding='utf-8') as f:
            for line in f:
                key, index = line.rstrip('\n').split('\t')
                ids[key] = int(index)
    return ids


def store_ids(file_path: str, ids: dict[str, int]) -> None:
    """ID 辞書を保存する。

    Args:
        file_path (str): ファイルパス
        ids (dict[str, int]): ID 辞書
    """
    with bz2.open(file_path, 'wt') as f:
        f.writelines(f'{key}\t{index}\n' for key, index in ids.items())


def create_id_dict(file_path: str) -> dict[str, int]:
    """ファイルから ID 辞書を作成する。

    Args:
        file_path (str): ファイルパス

    Returns:
        ID 辞書
    """
    id_dict = defaultdict(lambda: len(id_dict))
    # データパディング用
    id_dict[''] = 0
    if os.path.exists(file_path):
        id_dict.update(load_ids(file_path))
    return id_dict


def convert(
    input_file: str,
    x_file: str,
    y_file: str,
    word_ids: dict[str, int],
    pos_ids: dict[str, int]
):
    """文章データを形態素解析し、ID に変換して保存する。

    Args:
        input_file (str): 入力ファイルパス
        x_file (str): 訓練データ出力ファイルパス
        y_file (str): 教師データ出力ファイルパス
        word_ids (dict[str, int]): 単語 ID 辞書
        pos_ids (dict[str, int]): 品詞 ID 辞書
    """
    with bz2.open(input_file, 'rt', encoding='utf-8') as rf:
        tagger = MeCab.Tagger()
        # 訓練データ（単語 ID）
        x_train = []
        # 教師データ（品詞 ID）
        y_train = []

        for line in rf:
            # 単語 ID 配列
            words = []
            # 品詞 ID 配列
            poses = []
            # 文章を形態素解析する
            node = tagger.parseToNode(line)
            while node:
                # 単語
                word = node.surface
                # 品詞
                pos = node.feature.split(',')[0]
                # 単語を ID に変換して、配列に追加する
                words.append(word_ids[word])
                # 品詞を ID に変換して、配列に追加する
                poses.append(pos_ids[pos])
                # 次の単語にする
                node = node.next
            # 訓練データ、教師データを配列に追加する
            x_train.append(torch.tensor(words, dtype=torch.long))
            y_train.append(torch.tensor(poses, dtype=torch.long))

        # 訓練データにパディング処理をする
        x_train_padded = pad_sequence(x_train, batch_first=True, padding_value=0)
        # 訓練データをファイルに保存する
        np.save(x_file, x_train_padded.numpy())
        # 教師データにパディング処理をする
        y_train_padded = pad_sequence(y_train, batch_first=True, padding_value=-1)
        # 教師データをファイルに保存する
        np.save(y_file, y_train_padded.numpy())


def main():
    """
    文章データを形態素解析し、単語と品詞を ID に置換して、学習用データ（訓練データ、教師データ）を作成する。
    """
    # ID 辞書を読み込む
    word_ids = create_id_dict(WORD_FILE)
    pos_ids = create_id_dict(POS_FILE)

    # 学習用データを変換する
    print('学習用データ作成中...')
    convert(TRAIN_FILE, X_TRAIN_FILE, Y_TRAIN_FILE, word_ids, pos_ids)
    # テスト用データを変換する
    print('テスト用データ作成中...')
    convert(TEST_FILE, X_TEST_FILE, Y_TEST_FILE, word_ids, pos_ids)

    # ID 辞書を保存する
    store_ids(WORD_FILE, word_ids)
    store_ids(POS_FILE, pos_ids)
    print(f'単語 ID 数: {len(word_ids)}')
    print(f'品詞 ID 数: {len(pos_ids)}')


if __name__ == '__main__':
    main()

モデル学習を改善する

次に、モデル学習処理を効率化する。以下のような改善を行った。

バッチ処理化
学習率スケジューラ導入
TensorBoard 用ログ出力

初期の学習率は、エポック 10 程度で収束傾向になるように設定している。本来は、エポック数を増やし、初期の学習率を下げて学習を行うものなのかもしれない。

train_optim.py

import bz2
import logging

import numpy as np
import time
import torch
from torch.utils.data import TensorDataset, DataLoader
from torch.utils.tensorboard import SummaryWriter

from pos_tagger_lstm import PosTaggerLSTM

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

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

# 学習用訓練データ出力ファイル名
X_TRAIN_FILE = 'lstm-optim-x-train.npy'
# 学習用教師データ出力ファイル名
Y_TRAIN_FILE = 'lstm-optim-y-train.npy'

# バッチサイズ
BATCH_SIZE = 128
# 学習率
LEARNING_RATE = 1.0
# エポック数
EPOCHS = 10


# 学習率減衰の乗数
# GAMMA = 0.99


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

    # 単語 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)}

    # 訓練データファイルを読み込む
    x_train = np.load(X_TRAIN_FILE)
    x_train_tensor = torch.from_numpy(x_train)
    # 教師データファイルを読み込む
    y_train = np.load(Y_TRAIN_FILE)
    y_train_tensor = torch.from_numpy(y_train)

    # バッチ処理用にデータを準備する
    train_dataset = TensorDataset(x_train_tensor, y_train_tensor)
    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, pin_memory=True)

    print(f'word_ids: {len(word_ids)}, pos_ids: {len(pos_ids)}, x_train: {len(x_train)}, y_train: {len(y_train)}')

    # モデルに、単語 ID 数と、品詞 ID 数を指定する
    model = PosTaggerLSTM(len(word_ids) + 1, len(pos_ids), 100).to(device)
    # 最適化アルゴリズムに確率的勾配降下法を設定する
    optimizer = torch.optim.SGD(model.parameters(), lr=LEARNING_RATE)
    # 学習率スケジューラを設定する
    # 指定のエポック毎に、学習率に gamma を乗算した値を学習率に設定するスケジューラ
    # scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=GAMMA)
    # エポック毎に、学習率を指数関数的に更新する。初期学習率 * gamma ^ エポック数
    # scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=GAMMA)
    # 監視指標（scheduler.step 関数で渡す値）の変化率が threshold 値を下回ると、学習率に factor を欠けた値で学習率を更新する
    # threshold と比較する値(変化率)は、(前回の監視指標 - 現在の監視指標) / 前回の監視指標 となる
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.9, patience=0,
                                                           threshold=0.2, min_lr=0.01)
    # 損失関数に交差エントロピー損失を設定する
    # ignore_index=-1 で、教師データのパディング値である、-1 を無視する
    criterion = torch.nn.CrossEntropyLoss(ignore_index=-1)

    # TensorBoard ログ書き込み用
    writer = SummaryWriter()
    # 前回エポックの損失値（デバッグ表示用）
    epoch_loss_old = 0.00001

    for epoch in range(EPOCHS):
        epoch_start_time = time.time()
        # 損失の合計値
        loss_sum = 0.0
        for i, (x, y) in enumerate(train_loader):
            x = x.to(device)
            y = y.to(device)
            # 訓練する
            output = model(x)
            # 損失値(予測結果と正解データとの間の誤差)を計算する
            loss = criterion(output.view(-1, output.size(2)), y.view(-1))
            # 勾配を初期化する
            optimizer.zero_grad()
            # 勾配を計算する
            loss.backward()
            # 予測結果と正解データとの誤差 (損失) を小さくするように、モデルのパラメータを調整する
            optimizer.step()
            # 損失の合計値を更新する
            loss_sum += loss.item() * len(y)

            if i >= 1 and i % (1000 // BATCH_SIZE) == 0:
                print(f'{i * BATCH_SIZE}, Loss: {loss_sum / ((i + 1) * BATCH_SIZE):.4f}')

        epoch_loss = loss_sum / len(x_train)
        # 学習率を取得する
        lr = optimizer.param_groups[0]['lr']

        # TensorBoard 用ログを書き込む
        writer.add_scalar('Loss', loss_sum / len(x_train), epoch + 1)
        writer.add_scalar('Learning Rate', lr, epoch + 1)
        writer.flush()

        # 学習率を更新する（スケジューラによって引数が異なる）
        scheduler.step(epoch_loss)
        # この値が、スケジューラの threshold 値を下回ると、スケジューラが学習率が変更する
        print(f'rel_change: {(epoch_loss_old - epoch_loss) / epoch_loss_old}')
        epoch_loss_old = epoch_loss

        # モデルの保存
        outfile = f'lstm-optim-epoch{epoch + 1:02d}.model'
        torch.save({
            'epoch': epoch + 1,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'scheduler_state_dict': scheduler.state_dict(),
        }, outfile)
        elapsed_time = time.time() - epoch_start_time
        print(
            f'Epoch: {epoch + 1}, Loss: {loss_sum / len(x_train):.4f}, Learning Rate: {lr:.6f}, Time: {elapsed_time:.2f} s')


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

以下のような実行結果となった。バッチ化する前は、500 秒程度だったので、二割ほど高速化したことになる。 Colaboratory 上で、T4 GPU ランタイム上で動作確認すると、 600 秒ほどかかっていたものが、90 秒程度で完了するようになった。

Epoch: 1, Loss: 0.7359, Learning Rate: 1.000000, Time: 38.67 s
rel_change: 0.5249524916714136
Epoch: 2, Loss: 0.3496, Learning Rate: 1.000000, Time: 40.78 s
rel_change: 0.29967549745495214
Epoch: 3, Loss: 0.2448, Learning Rate: 1.000000, Time: 38.42 s
rel_change: 0.1830075276367609
Epoch: 4, Loss: 0.2000, Learning Rate: 1.000000, Time: 39.26 s
rel_change: 0.13048517080262004
Epoch: 5, Loss: 0.1739, Learning Rate: 0.900000, Time: 40.58 s
rel_change: 0.0957767101716502
Epoch: 6, Loss: 0.1573, Learning Rate: 0.900000, Time: 40.53 s
rel_change: 0.07944206187798325
Epoch: 7, Loss: 0.1448, Learning Rate: 0.810000, Time: 37.85 s
rel_change: 0.06455950856455013
Epoch: 8, Loss: 0.1354, Learning Rate: 0.729000, Time: 38.40 s
rel_change: 0.05247316818668716
Epoch: 9, Loss: 0.1283, Learning Rate: 0.729000, Time: 37.97 s
rel_change: 0.0500079453389081
Epoch: 10, Loss: 0.1219, Learning Rate: 0.656100, Time: 39.10 s
経過時間: 394.22 s

学習率の確認

学習処理を実行すると、TensorBoard 用のログデータが、./runs ディレクトリに出力される。 TensorBoard を起動して、ログを確認する方法は、以下のようにする。 runs ディレクトの親ディレクトリで、以下を実行する

tensorboard --logdir=./runs

TensorBoard が起動したら、ブラウザで、http://localhost:6006/ にアクセスする。 Colaboratory 上であれば、以下のコードを実行すると、Colaboratory 上に TensorBoard の画面が表示される。

%load_ext tensorboard
%tensorboard --logdir=./runs

推論をする

学習済みのモデルを利用して、推論をする。推論には、テスト用データを利用する。

infer_optim.py

import bz2
import logging

import numpy as np
import time
import torch
from torch.utils.data import TensorDataset, DataLoader

from pos_tagger_lstm import PosTaggerLSTM

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

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

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

# モデルファイル名
MODEL_FILE = 'lstm-optim-epoch10.model'

# バッチサイズ
BATCH_SIZE = 128


def main():
    """
    LSTM モデル推論を行う
    """
    print('Loading datas...')
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

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

    # テスト用訓練データファイルを読み込む
    x_test = np.load(X_TEST_FILE)
    x_test_tensor = torch.from_numpy(x_test)
    # テスト用教師データファイルを読み込む
    y_test = np.load(Y_TEST_FILE)
    y_test_tensor = torch.from_numpy(y_test)

    train_dataset = TensorDataset(x_test_tensor, y_test_tensor)
    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=False, pin_memory=True)

    model = PosTaggerLSTM(len(word_ids) + 1, len(pos_ids), 100).to(device)
    checkpoint = torch.load(MODEL_FILE, weights_only=True, map_location=device)
    model.load_state_dict(checkpoint['model_state_dict'])

    data_count = 0
    model.eval()
    print('Start inference...')
    with torch.no_grad():
        start_time = time.time()
        ok = 0
        for x, y in train_loader:
            x = x.to(device)
            y = y.to(device)
            output = model(x)
            # 推論結果から、回答を取り出す
            ans = torch.argmax(output, dim=2)
            # 教師データから、パディング値を除外するための配列を作成する(パディング値なら False となる配列)
            mask = y != -1
            # 回答から、正解でかつ、パディングデータでないものを抽出する
            # * mask の部分は、AND 条件と判定される。
            ok += ((ans == y) * mask).sum().item()
            # mask の値が、True の項目合計数を求める
            data_count += mask.sum().item()
        print(f'OK: {ok}, Total: {data_count}, Accuracy: {ok / data_count:.4f}')
        elapsed_time = time.time() - start_time
        print(f'経過時間: {elapsed_time:.2f} s ({data_count / elapsed_time:.0f} times/s)')


if __name__ == '__main__':
    main()

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

Loading datas...
Start inference...
OK: 424164, Total: 440767, Accuracy: 0.9623
経過時間: 4.28 s (102958 times/s)

まとめ

今回のバッチ化対応で、CPU や、GPU で速度改善ができた。次の記事では、LSTM モデルを多層にしたり、双方向にする体験をしていこうと思う。

前の記事：自然言語処理編 - LSTM モデル学習その弐
次の記事：自然言語処理編 - LSTM モデル学習その肆

いちりのつれづれ

このブログを検索

AI 物見遊山 - 自然言語処理編 - LSTM モデル学習その参

自然言語処理編 - LSTM モデル学習その参

学習用データ生成処理を変更する

モデル学習を改善する

学習率の確認

推論をする

まとめ

AI 物見遊山 - 自然言語処理編 - LSTM モデル学習 その参

自然言語処理編 - LSTM モデル学習 その参

学習用データ生成処理を変更する

モデル学習を改善する

学習率の確認

推論をする

まとめ

AI 物見遊山 - 自然言語処理編 - LSTM モデル学習その参

自然言語処理編 - LSTM モデル学習その参