Optunaを用いた文書分類モデルのハイパラ最適化

はじめに

より良い機械学習モデルの構築のために, Batch sizeやDropout率といったハイパーパラメータ（ハイパラ）の調整は大きな課題の1つです.
本記事では, ニュース記事のカテゴリを分類する文書分類モデルのハイパラ最適化について解説します. 具体的には, モデルとして事前学習済みのBERTを使用し, ファインチューニング時のハイパラ最適化方法を扱います.
ハイパラ最適化には, 機械学習モデルのハイパーパラメータ自動最適化フレームワークの1つであるOptunaを使用します.

実験環境

本記事の実験環境, ソースコードは下記のNotebookで公開しています.

• Optunaで文書分類モデルのハイパラチューニング (Google Colab Notebook)

参考記事

本記事で実装した文書分類モデルやOptunaを用いたハイパラ最適化のコードは下記の記事を参考にさせていただきました.

• 【PyTorch】BERTを用いた日本語文書分類入門
• Optunaでハイパーパラメータの自動チューニング -Pytorch Lightning編-
• optuna/optuna-examples (github)
• PyTorch Lightningによる多クラス分類の実装

Livedoorニュースコーパスの取得

Livedoorニュースコーパスは「livedoorニュース」の下記9種類のニュース記事からなるコーパスです.

• トピックニュース
• Sports Watch
• ITライフハック
• 家電チャンネル
• MOVIE ENTER
• 独女通信
• エスマックス
• livedoor HOMME
• Peachy

本記事では, Livedoorニュースコーパスを用いて、ある記事のニュースカテゴリ（上記9種類のうちいずれか）を分類する文書分類モデルを構築します.
こちらのサイトで配布されているため, ダウンロードおよび前処理を行います.

# livedoorニュースコーパスのダウンロード
!wget https://www.rondhuit.com/download/ldcc-20140209.tar.gz
!tar zxvf ldcc-20140209.tar.gz

# 整形結果格納用ファイル作成
!echo -e "filename\tarticle"$(for category in $(basename -a `find ./text -type d` | grep -v text | sort); do echo -n "\t"; echo -n $category; done) > ./text/livedoor.tsv

# カテゴリごとに格納
!for filename in `basename -a ./text/dokujo-tsushin/dokujo-tsushin-*`; do echo -n "$filename"; echo -ne "\t"; echo -n `sed -e '1,3d' ./text/dokujo-tsushin/$filename`; echo -e "\t1\t0\t0\t0\t0\t0\t0\t0\t0"; done >> ./text/livedoor.tsv
!for filename in `basename -a ./text/it-life-hack/it-life-hack-*`; do echo -n "$filename"; echo -ne "\t"; echo -n `sed -e '1,3d' ./text/it-life-hack/$filename`; echo -e "\t0\t1\t0\t0\t0\t0\t0\t0\t0"; done >> ./text/livedoor.tsv
!for filename in `basename -a ./text/kaden-channel/kaden-channel-*`; do echo -n "$filename"; echo -ne "\t"; echo -n `sed -e '1,3d' ./text/kaden-channel/$filename`; echo -e "\t0\t0\t1\t0\t0\t0\t0\t0\t0"; done >> ./text/livedoor.tsv
!for filename in `basename -a ./text/livedoor-homme/livedoor-homme-*`; do echo -n "$filename"; echo -ne "\t"; echo -n `sed -e '1,3d' ./text/livedoor-homme/$filename`; echo -e "\t0\t0\t0\t1\t0\t0\t0\t0\t0"; done >> ./text/livedoor.tsv
!for filename in `basename -a ./text/movie-enter/movie-enter-*`; do echo -n "$filename"; echo -ne "\t"; echo -n `sed -e '1,3d' ./text/movie-enter/$filename`; echo -e "\t0\t0\t0\t0\t1\t0\t0\t0\t0"; done >> ./text/livedoor.tsv
!for filename in `basename -a ./text/peachy/peachy-*`; do echo -n "$filename"; echo -ne "\t"; echo -n `sed -e '1,3d' ./text/peachy/$filename`; echo -e "\t0\t0\t0\t0\t0\t1\t0\t0\t0"; done >> ./text/livedoor.tsv
!for filename in `basename -a ./text/smax/smax-*`; do echo -n "$filename"; echo -ne "\t"; echo -n `sed -e '1,3d' ./text/smax/$filename`; echo -e "\t0\t0\t0\t0\t0\t0\t1\t0\t0"; done >> ./text/livedoor.tsv
!for filename in `basename -a ./text/sports-watch/sports-watch-*`; do echo -n "$filename"; echo -ne "\t"; echo -n `sed -e '1,3d' ./text/sports-watch/$filename`; echo -e "\t0\t0\t0\t0\t0\t0\t0\t1\t0"; done >> ./text/livedoor.tsv
!for filename in `basename -a ./text/topic-news/topic-news-*`; do echo -n "$filename"; echo -ne "\t"; echo -n `sed -e '1,3d' ./text/topic-news/$filename`; echo -e "\t0\t0\t0\t0\t0\t0\t0\t0\t1"; done >> ./text/livedoor.tsv

# ファイルの内容を確認
!head -10 ./text/livedoor.tsv
(省略)

文書分類モデルの構築

今回は, PyTorchの軽量ラッパーであるPyTorch Lightningを使って文書分類モデルを実装していきます.
PyTorch Lightningを使うことで, PyTorchを用いたモデルの実装において煩雑になりがちなtrain/testループ等をシンプルに書くことができます.

PyTorch Ligntningの詳細はこちらの記事が参考になります.

モデル構築の流れ

PyTorch Lightningを使ったモデル構築の流れはおおまかに以下の3つです.
LightningModuleとは, torch.nn.Moduleを拡張したようなクラスになっていて, モデルの定義だけでなく, lossの計算やoptimizerの定義などをまとめることができます.

• データの読み込み
• Datasetの定義
• LightningModuleの定義

データの読み込み

まずはさきほどダウンロード・前処理したlivedoorニュースコーパスをDataFrameに読み込みます.

import pandas as pd
from sklearn.model_selection import train_test_split
from tabulate import tabulate

# データの読込
df = pd.read_csv('./text/livedoor.tsv', sep='\t')

# データの分割
categories = ['dokujo-tsushin', 'it-life-hack', 'kaden-channel', 'livedoor-homme', 'movie-enter', 'peachy', 'smax', 'sports-watch', 'topic-news']
train, valid_test = train_test_split(df, test_size=0.2, shuffle=True, random_state=123, stratify=df[categories])
valid, test = train_test_split(valid_test, test_size=0.5, shuffle=True, random_state=123, stratify=valid_test[categories])
train.reset_index(drop=True, inplace=True)
valid.reset_index(drop=True, inplace=True)
test.reset_index(drop=True, inplace=True)

# 事例数の確認
table = [['train'] + [train[category].sum() for category in categories],
         ['valid'] + [valid[category].sum() for category in categories],
         ['test'] + [test[category].sum() for category in categories]]
headers = ['data'] + categories
print(tabulate(table, headers, tablefmt='grid'))

Datasetの定義

学習・検証時に使用するデータセットを定義します.
具体的には, PyTorchの Dataset クラスを継承させた MyDataset クラスを作ります.
__getitem__には, indexを引数として, データを返す処理を記述します.

また, 作成した MyDataset クラス内で テキストデータのトークンid化やPaddingといった前処理を行うことができます.

# Datasetの定義
class MyDataset(Dataset):
    def __init__(self, X, y, tokenizer, max_len):
        self.X = X
        self.y = y
        self.tokenizer = tokenizer
        self.max_len = max_len

    def __len__(self):  # len(Dataset)で返す値を指定
        return len(self.y)

    def __getitem__(self, index):  # Dataset[index]で返す値を指定
        text = self.X[index]
        inputs = self.tokenizer.encode_plus(
            text,
            add_special_tokens=True,
            max_length=self.max_len,
            pad_to_max_length=True
        )
        ids = inputs['input_ids']
        mask = inputs['attention_mask']

        return {
            'ids': torch.LongTensor(ids),
            'mask': torch.LongTensor(mask),
            'labels': torch.Tensor(self.y[index])
        }
    
    @staticmethod
    def collate_fn(batch):
        ids_batch = pad_sequence([item["ids"] for item in batch], padding_value=0).transpose(0, 1)
        mask_batch = pad_sequence([item["mask"] for item in batch], padding_value=0).transpose(0, 1)
        labels_batch = pad_sequence([item["labels"] for item in batch]).transpose(0, 1)
        return ids_batch, mask_batch, labels_batch

LightningModule(モデル)の定義

次に, LightningModule（モデル）を定義するために, pl.LightningModule を継承したクラスを作成します. PyTorch Lightningでは, LightningModule内にネットワークやloss関数, optimizerを定義します.
今回は, BERTを用いた文書分類モデルを作成するため, BERTClassifier というクラス名のLightnignModuleを作成しました.

pl.LightningModuleでは様々なメソッドが定義されており, 例えば train_epoch_endやvalidation_epoch_endといったメソッドをオーバーライドすることで, 学習または検証ループのエポック終了ごとに実行する処理を記述することができます.
他にもさまざまなメソッドが定義されていますので詳細は公式ドキュメントを参照して下さい.

class BERTClassifier(pl.LightningModule):
    def __init__(self, pretrained, drop_rate, output_size, lr):
        super().__init__()
        self.bert = BertModel.from_pretrained(pretrained)
        self.drop = torch.nn.Dropout(drop_rate)
        self.fc = torch.nn.Linear(768, output_size)  # BERTの出力に合わせて768次元を指定
        self.lr = lr
        self.criterion = torch.nn.BCEWithLogitsLoss()

    def forward(self, ids, mask):
        outputs = self.bert(ids, attention_mask=mask, return_dict=True)
        out = self.fc(self.drop(outputs["last_hidden_state"][:, 0, :]))
        return out
        
    def training_step(self, batch, batch_idx):
        ids_batch = batch[0]
        mask_batch = batch[1]
        labels_batch = batch[2]        
        out = self(ids_batch, mask_batch)
        loss = self.criterion(out, labels_batch)
        self.log('train_loss', loss)        
        pred_labels = out.argmax(dim=1).detach().cpu().tolist()
        tgt_labels = labels_batch.argmax(dim=1).detach().cpu().tolist()
        return {
            "loss": loss, 
            "pred_labels": pred_labels,
            "tgt_labels": tgt_labels
        }

    def train_epoch_end(self, train_step_outputs):
        pred_labels = []
        tgt_labels = []
        for out in train_step_outputs:
            pred_labels.extend(out["pred_labels"])
            tgt_labels.extend(out["tgt_labels"])
        train_f1 = f1_score(tgt_labels, pred_labels, average="macro")
        self.log("train_f1", train_f1)

    def validation_step(self, batch, batch_idx):
        ids_batch = batch[0]
        mask_batch = batch[1]
        labels_batch = batch[2]        
        out = self(ids_batch, mask_batch)
        loss = self.criterion(out, labels_batch)
        self.log('valid_loss', loss)
        pred_labels = out.argmax(dim=1).detach().cpu().tolist()
        tgt_labels = labels_batch.argmax(dim=1).detach().cpu().tolist()
        return {
            "loss": loss, 
            "pred_labels": pred_labels,
            "tgt_labels": tgt_labels
        }

    def validation_epoch_end(self, validation_step_outputs):
        pred_labels = []
        tgt_labels = []
        for out in validation_step_outputs:
            pred_labels.extend(out["pred_labels"])
            tgt_labels.extend(out["tgt_labels"])
        val_f1 = f1_score(tgt_labels, pred_labels, average="macro")
        self.log("val_f1", val_f1)
        
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=self.lr)
        return optimizer

モデルの学習/検証

これまで定義した MyDataset からDataloderを作成し, BERTClassifierの学習します.
PyTorch Lightningでは, 下記のように trainer を作成し, trainer.fit で定義したモデルの学習・検証ループを実行することができます.

MAX_LEN = 128
batch_size = 32
epochs = 5

# 事前学習済みモデルの指定
pretrained = 'cl-tohoku/bert-base-japanese-whole-word-masking'

# tokenizerの取得
tokenizer = BertJapaneseTokenizer.from_pretrained(pretrained)

# modelの作成
model = BERTClassifier(pretrained, drop_rate=0.1, output_size=9, lr=2e-5)

# Datasetの作成
dataset_train = MyDataset(train['article'], train[categories].values, tokenizer, MAX_LEN)
dataset_valid = MyDataset(valid['article'], valid[categories].values, tokenizer, MAX_LEN)

# Dataloaderの作成
dataloader_train = DataLoader(dataset_train, batch_size=batch_size, shuffle=True, collate_fn=dataset_train.collate_fn)
dataloader_valid = DataLoader(dataset_valid, batch_size=batch_size, shuffle=False, collate_fn=dataset_valid.collate_fn)

trainer = pl.Trainer(
    logger=True,
    max_epochs=epochs,
    checkpoint_callback=False,
    gpus=1,
)

trainer.fit(model, dataloader_train, dataloader_valid)

これで, 文書分類モデルの学習・検証ループまでを実装することができました.
それでは, 次にこのモデルに対するハイパラ最適化を行います.

ハイパラ最適化

ハイパラ最適化の流れ

さきほど作成したモデルの学習/検証のコードを参考に, Optunaによるハイパラ最適化を行います.
Optunaでは主に下記2ステップによりハイパラ最適化を行います.

まず, 目的関数の定義では, 最適化したいスコア（loss, 精度 等）を返す objective 関数を作成します.
次に最適化の実行により, lossであれば最小化, 精度であれば最大化を行います. 　　

今回実装したコードでは, 検証データ（valid）に対する分類精度(F1)を最大化するようにハイパラ最適化を行います.

• 目的関数（objective）の定義
• 最適化の実行

目的関数の定義

こちらが作成した目的関数（objective）です. 今回調整する対象のハイパーパラメータは, 学習率(lr)とドロップアウト確率（drop_rate）としました.
Optunaでは, trial.suggest_floatのように調整したいパラメータとその範囲を指定することができます.

objectiveでは, ハイパラ最適化により最大化または最小化させたいスコアを返します.
今回は, 検証データ（valid）に対する分類精度(F1)を最大化するようにハイパラ最適化を行います.

def objective(trial):
    MAX_LEN = 128
    batch_size = 32
    epochs = 5

    # 学習率（lr）とドロップアウト確率（drop_rate）を最適化
    lr = trial.suggest_float("lr", 2e-5, 2e-4)
    drop_rate = trial.suggest_float("drop_rate", 0.1, 0.5)

    # 事前学習済みモデルの指定
    pretrained = 'cl-tohoku/bert-base-japanese-whole-word-masking'

    # tokenizerの取得
    tokenizer = BertJapaneseTokenizer.from_pretrained(pretrained)

    # modelの作成
    model = BERTClassifier(pretrained, drop_rate=drop_rate, output_size=9, lr=lr)

    # Datasetの作成
    dataset_train = MyDataset(train['article'], train[categories].values, tokenizer, MAX_LEN)
    dataset_valid = MyDataset(valid['article'], valid[categories].values, tokenizer, MAX_LEN)

    # Dataloaderの作成
    dataloader_train = DataLoader(dataset_train, batch_size=batch_size, shuffle=True, collate_fn=dataset_train.collate_fn)
    dataloader_valid = DataLoader(dataset_valid, batch_size=batch_size, shuffle=False, collate_fn=dataset_valid.collate_fn)

    trainer = pl.Trainer(
        logger=True,
        max_epochs=epochs,
        checkpoint_callback=False,
        gpus=1,
        callbacks=[PyTorchLightningPruningCallback(trial, monitor="val_f1")]
    )

    hyperparameters = dict(lr=lr, drop_rate=drop_rate)
    trainer.logger.log_hyperparams(hyperparameters)

    trainer.fit(model, dataloader_train, dataloader_valid)
    return trainer.callback_metrics["val_f1"].item() # val_f1を（最適化）最大化する

最適化の実行

最後に定義した目的関数を用いて, Studyにより最適化（スコア最大化）を実行します.
optuna.create_studyをdirection="maximize"とすることで, スコア最大化をすることができます.

pruner = optuna.pruners.MedianPruner()は, ハイパラ最適化における枝刈りの基準を表しています.
MedianPrunerでは, その名の通り, 過去の試行（trial）の同じepochにおける値と比較して, それらの中央値よりもスコアが悪ければ試行を打ち切る基準となっています.

今回のコードでは, n_trials（試行回数）を10に設定しました.
これにより, 今回実装した文書分類モデルのファインチューニングが10回実行されます.

pruner = optuna.pruners.MedianPruner()

study = optuna.create_study(direction="maximize", pruner=pruner)
study.optimize(objective, n_trials=10)

print("Number of finished trials: {}".format(len(study.trials)))

print("Best trial:")
trial = study.best_trial

print("  Value: {}".format(trial.value))

print("  Params: ")
for key, value in trial.params.items():
    print("    {}: {}".format(key, value))

ハイパラ最適化が完了すると下記のように結果を表示します.
Valueは, 試行回数10回のうち, 最もスコアが良かった検証データ（valid）に対するF1スコアを表しています. また, その際のlrとdrop_rateも確認することができます.

Number of finished trials: 10
Best trial:
  Value: 0.9065759778022766
  Params: 
    lr: 2.4313051844385965e-05
    drop_rate: 0.34326517678045065

おわりに

ニュース記事のカテゴリを分類する文書分類モデルのハイパラ最適化について解説しました. 文書分類モデルとしてBERTを使用し, 分類タスクのファインチューニング時のハイパラ最適化を行いました.

PyTorchやPytorch Lightningでモデルを実装する際には, はじめからOptunaのobjective関数を意識してコーディングすることで比較的容易にハイパラ最適化の導入が実現できそうです.