OptunaでFlair NERのパラメーターチューニング
とある仕事で，固有表現抽出器の性能をぎりぎりまで向上させる必要があったため，自作のコーパスに対して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