OptunaでFlair NERのパラメーターチューニング

1月 24, 2021

とある仕事で，固有表現抽出器の性能をぎりぎりまで向上させる必要があったため，自作のコーパスに対してFlairでNERを学習し，Optunaでハイパーパラメータチューニングを行なった．本記事ではそのときのコードを示す．

Embeddingについては事前にどれが有効かわかっていたので固定にし，Learning rateなど細かい設定をチューニングするつくりにした．

学習部の関数(Objective)では，micro F1スコアを返すようになっているため，optuna.create_study(direction=’maximize’)と，directionにmaximizeを設定．

評価関数(スクリプト内ではoptimizerと定義)もチューニングしようとするが，デフォルトのもの以外で性能がガタ落ちしたため，取りやめ．私の設定が甘かったのか？

また，optunaのパラメーター設定で必要な関数は，私の場合，trial.suggest_intとtrial.suggest_floatとtrial.suggest_categoricalの3つだけで十分だと感じた．設定の方法は下記スクリプトに記載しているのでどうぞ参考にされたし．

本記事の最下部では，このチューニング用スクリプトを実行して出力されるcsvの例を示す．

flair_ner_optimization.py

# このスクリプトはflair==0.4.3で動作確認
from pathlib import Path
from argparse import ArgumentParser
from flair.models import SequenceTagger
from flair.data import Dictionary, Corpus
from flair.datasets import ColumnCorpus
from flair.embeddings import (
    # TransformerWordEmbeddings,
    StackedEmbeddings,
    BytePairEmbeddings,
    CharacterEmbeddings,
    ELMoEmbeddings,
)
from flair.trainers import ModelTrainer
# import allennlp
import torch
import optuna


def objective(trial):
    
    parser = ArgumentParser()

    parser.add_argument("corpus", type=Path)
    parser.add_argument("output", type=Path)
    parser.add_argument("--transformer", type=str)
    parser.add_argument("--fine_tune", action="store_true")
    parser.add_argument("--device", type=str, default="cpu")
    parser.add_argument("--elmo", type=str)
    parser.add_argument("--bpe", action='store_true')
    parser.add_argument("--char", action='store_true')
    parser.add_argument("--electra_small", action='store_true')
    parser.add_argument("--electra_base", action='store_true')
    parser.add_argument("--electra_large", action='store_true')

    args = parser.parse_args()
    
    lr = trial.suggest_float('lr', 0.05, 0.3, step=0.05)
    dropout = trial.suggest_float('dropout', 0.3, 0.6, step=0.1)
    locked_dropout = trial.suggest_float('locked_dropout', 0.3, 0.6, step=0.1)
    hidden_size = trial.suggest_categorical('hidden_size', [32, 64, 128, 256])
    mini_batch_size = trial.suggest_categorical('mini_batch_size', [16, 32, 64])
    rnn_layers = trial.suggest_int('rnn_layers', 1, 3, step=1)
    # anneal_factor = trial.suggest_float('anneal_factor', 0.1, 0.6, step=0.1)
    word_dropout= trial.suggest_float('word_dropout', 0.05, 0.15, step=0.05)
    beta = trial.suggest_float('beta', 0.7, 1.0, step=0.1)
    weight_decay = trial.suggest_float('weight_decay', 0.0001, 0.0005, step=0.0001)
    # optimizer = trial.suggest_categorical('optimizer', ['Default', 'AdamP', 'AdaBelief', 'DiffGrad', 'SGDP'])
    
    corpus_dir = args.corpus
    output_dir = args.output
    labels_txt = args.corpus / "labels.txt"
    fine_tune = args.fine_tune
    device = args.device
    transformer = args.transformer
    elmo = args.elmo
    char = args.char
    bpe = args.bpe
    electra_small = args.electra_small
    electra_base = args.electra_base
    electra_large = args.electra_large

    # embeddings
    embedding_types = []
    if transformer:
        embedding_types.append(TransformerWordEmbeddings(transformer, from_tf=Path(transformer).exists(), fine_tune=fine_tune))
    if elmo:
        # ELMoをEmbeddingに使う場合は，指定するパス（フォルダ）に，elmo_options.jsonとelmo_weights.hdf5が入っている必要がある．
        embedding_types.append(ELMoEmbeddings("custom",options_file=str(Path(elmo)/'elmo_options.json'), weight_file=str(Path(elmo)/'elmo_weights.hdf5')))
    if bpe:
        embedding_types.append(BytePairEmbeddings("en"))
    if char:
        embedding_types.append(CharacterEmbeddings())
    if electra_small:
        embedding_types.append(TransformerWordEmbeddings("google/electra-small-discriminator", fine_tune=fine_tune))
    if electra_base:
        embedding_types.append(TransformerWordEmbeddings("google/electra-base-discriminator", fine_tune=fine_tune))
    if electra_large:
        embedding_types.append(TransformerWordEmbeddings("google/electra-large-discriminator", fine_tune=fine_tune))
    embeddings = StackedEmbeddings(embedding_types)

    # prepare corpus
    columns = {0: "text", 1: "ner"}
    corpus = ColumnCorpus(
        corpus_dir,
        columns,
        train_file="train.tsv",
        test_file="test.tsv",
        dev_file="devel.tsv",
        # column_delimiter="\t"
        # corpus_dir, columns, train_file="dev.txt", test_file="test.txt", dev_file="dev.txt", column_delimiter="\t"
    )

    # create label dictionary
    # labels = labels_txt.read_text().split("\n")
    # tag_dictionary = Dictionary()
    # for lbl in labels:
    #     tag_dictionary.add_item(lbl)
    tag_type = "ner"
    tag_dictionary = corpus.make_tag_dictionary(tag_type=tag_type)
    print(tag_dictionary)

    # define model
    # tagger = SequenceTagger(hidden_size=512, embeddings=embeddings, tag_dictionary=tag_dictionary, tag_type=tag_type)
    tagger = SequenceTagger(
        hidden_size=hidden_size,
        rnn_layers=rnn_layers,
        embeddings=embeddings,
        tag_dictionary=tag_dictionary,
        dropout=dropout,
        locked_dropout=locked_dropout,
        word_dropout=word_dropout,
        beta=beta,
        tag_type=tag_type, use_crf=True
    )
    # 参考のため， flairのリポジトリからSequenceTaggerの引数を取得
    # def __init__(
    #         self,
    #         hidden_size: int,
    #         embeddings: TokenEmbeddings,
    #         tag_dictionary: Dictionary,
    #         tag_type: str,
    #         use_crf: bool = True,
    #         use_rnn: bool = True,
    #         rnn_layers: int = 1,
    #         dropout: float = 0.0,
    #         word_dropout: float = 0.05,
    #         locked_dropout: float = 0.5,
    #         reproject_embeddings: Union[bool, int] = True,
    #         train_initial_hidden_state: bool = False,
    #         rnn_type: str = "LSTM",
    #         pickle_module: str = "pickle",
    #         beta: float = 1.0,
    #         loss_weights: Dict[str, float] = None,
    # ):

    # if optimizer == 'AdamP':
    #     from torch_optimizer import AdamP as opt
    # elif optimizer == 'AdaBelief':
    #     from torch_optimizer import AdaBelief as opt
    # elif optimizer == 'DiffGrad':
    #     from torch_optimizer import DiffGrad as opt
    # elif optimizer == 'SGDP':
    #     from torch_optimizer import SGDP as opt
        
    # if optimizer == 'Default':
    #     trainer: ModelTrainer = ModelTrainer(tagger, corpus)
    # else:
    #     trainer: ModelTrainer = ModelTrainer(tagger, corpus, optimizer=opt)
        
    trainer: ModelTrainer = ModelTrainer(tagger, corpus)
        
    result = trainer.train(
        base_path=output_dir,
        learning_rate=lr,
        mini_batch_size=mini_batch_size,
        max_epochs=25,  #
        patience=5,
        monitor_test=True,
        embeddings_storage_mode=device,
        save_final_model=False, # パラメータ探索が目的のため
        # anneal_factor=anneal_factor,
        weight_decay=weight_decay,
    )

    # 参考のため， flairのリポジトリからtrainの引数を取得
    # TRAINING
    # trainer = ModelTrainer(tagger, corpus, use_tensorboard=True, optimizer=torch.optim.AdamW)
    # trainer.train(
    #     base_path=output_dir,
    #     learning_rate=1e-3,
    #     mini_batch_size=16,
    #     mini_batch_chunk_size=None,
    #     max_epochs=100,
    #     # scheduler=AnnealOnPlateau,
    #     # cycle_momentum=False,
    #     # anneal_factor=0.5,
    #     patience=3,
    #     initial_extra_patience=0,
    #     min_learning_rate=1e-8,
    #     train_with_dev=False,
    #     monitor_train=False,
    #     monitor_test=False,
    #     embeddings_storage_mode=device,
    #     checkpoint=True,
    #     save_final_model=True,
    #     # anneal_with_restarts=False,
    #     # anneal_with_prestarts=False,
    #     # batch_growth_annealing=False,
    #     shuffle=True,
    #     param_selection_mode=False,
    #     write_weights=False,
    #     num_workers=6,
    #     sampler=None,
    #     use_amp=False,
    #     amp_opt_level="O1",
    #     eval_on_train_fraction=0.0,
    #     eval_on_train_shuffle=False,
    # )
    
    return result['test_score']


if __name__ == "__main__":
    
    study = optuna.create_study(direction='maximize')
    study.optimize(objective, n_trials=50)
    
    print('params:', study.best_params)
    print('value:', study.best_value)
    
    study_df = study.trials_dataframe()
    study_df.to_csv('./flair_ner_optimization.csv')

上のスクリプトを実行し，得られる結果(csv)．

これをみると，各ステップで探索されたパラメーターが一目瞭然である．結果の中から最もF1スコア(ここでいうvalue)の高いパラメータの組み合わせを本番の学習で設定することによって，性能改善の効果が期待できる．

筆者の場合は，0.2~0.4くらいしか性能が上がっていないが，チューニング前よりは確かに上がっているため，やらないよりはましだろう．また，深層学習ベースの学習器の場合，ある1つの設定でたまたま高スコアがでることがよくあるため，こうして何度も学習を行い，どのステップにおいても同じくらいのスコアがでるのを確認することは，モデルのデバッグを行うことにおいても有効なのではないかと思う．

	number	value	datetime_start	datetime_complete	duration	params_beta	params_dropout	params_hidden_size	params_locked_dropout	params_lr	params_mini_batch_size	params_rnn_layers	params_weight_decay	params_word_dropout	state
0	0	0.763867827	10:59.9	37:11.1	0 days 00:26:11.182141000	0.8	0.6	256	0.6	0.2	32	1	0.0005	0.15	COMPLETE
1	1	0.800172265	37:11.1	14:16.9	0 days 00:37:05.824783000	1	0.4	256	0.3	0.25	16	2	0.0003	0.15	COMPLETE
2	2	0.793049793	14:16.9	34:33.8	0 days 00:20:16.894379000	1	0.4	128	0.3	0.2	64	1	0.0002	0.15	COMPLETE
3	3	0.774522723	34:33.8	55:13.7	0 days 00:20:39.848606000	0.7	0.4	64	0.4	0.15	64	2	0.0003	0.15	COMPLETE
4	4	0.759173047	55:13.7	15:23.5	0 days 00:20:09.795247000	0.8	0.6	64	0.4	0.05	64	1	0.0002	0.15	COMPLETE
5	5	0.724664384	15:23.5	41:44.8	0 days 00:26:21.303128000	0.7	0.6	32	0.4	0.15	32	2	0.0002	0.1	COMPLETE
6	6	0.786737225	41:44.8	08:52.6	0 days 00:27:07.824726000	0.8	0.3	64	0.4	0.2	32	3	0.0002	0.15	COMPLETE
7	7	0.783839779	08:52.6	44:54.4	0 days 00:36:01.805369000	1	0.6	256	0.4	0.2	16	1	0.0004	0.15	COMPLETE
8	8	0.773110651	44:54.4	05:12.4	0 days 00:20:17.958787000	0.7	0.4	128	0.4	0.05	64	1	0.0005	0.15	COMPLETE
9	9	0.76408377	05:12.4	41:33.8	0 days 00:36:21.409957000	0.8	0.5	64	0.6	0.1	16	2	0.0005	0.05	COMPLETE