Hugging FaceでEarlyStoppingを使用する方法

こんちには。

データアナリティクス事業本部 機械学習チームの中村です。

Hugging Faceのライブラリの使い方紹介記事第２弾です。

今回は、学習時にEarly Stoppingを使用する方法をご紹介します。

Early Stoppingとは

Deep Learningなどのモデルを学習する際には、通常trainデータとvalidデータ(検定データ)に分割し、 trainデータでモデル自体を更新し、validデータで学習外データに対する汎化性能を確認します。

より具体的にはvalidデータの損失関数などの指標を確認し、最良となるようなepochで学習を停止するなどの操作をします。 これがEarly Stoppingと一般的に呼ばれるものです。

実行環境

今回はGoogle Colaboratory環境で実行しました。

ハードウェアなどの情報は以下の通りです。

• GPU: Tesla P100 (GPUメモリ16GB搭載)
• CUDA: 11.1
• メモリ: 13GB

主なライブラリのバージョンは以下となります。

• transformers: 4.22.1
• datasets: 2.4.0

インストール

transformersとdatasetsをインストールします。

!pip install transformers datasets

また事前学習モデルの依存モジュールをインストールします。

!pip install fugashi
!pip install ipadic
!pip install sentencepiece

ベースとするコード

今回のベースとするコードは以下のとおりです。

from datasets import load_dataset
from transformers import AutoTokenizer
from transformers import AutoModelForSequenceClassification
from transformers import TrainingArguments
from transformers import Trainer
from sklearn.metrics import accuracy_score, f1_score
import torch

# データセットのロード
dataset = load_dataset("tyqiangz/multilingual-sentiments", "japanese")

# # 実験のためデータセットを縮小したい場合はコチラを有効化
# dataset = DatasetDict({
#     "train": dataset['train'].select(range(1000)),
#     "validation": dataset['validation'].select(range(1000)),
#     "test": dataset['test'].select(range(1000)),
# })

# トークナイザのロード
model_ckpt = "cl-tohoku/bert-base-japanese-whole-word-masking"
tokenizer = AutoTokenizer.from_pretrained(model_ckpt)

# トークナイズ処理
def tokenize(batch):
    return tokenizer(batch["text"], padding=True, truncation=True)
dataset_encoded = dataset.map(tokenize, batched=True, batch_size=None)

# 事前学習モデルのロード
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
num_labels = 3
model = (AutoModelForSequenceClassification
    .from_pretrained(model_ckpt, num_labels=num_labels)
    .to(device))

# メトリクスの定義
def compute_metrics(pred):
    labels = pred.label_ids
    preds = pred.predictions.argmax(-1)
    f1 = f1_score(labels, preds, average="weighted")
    acc = accuracy_score(labels, preds)
    return {"accuracy": acc, "f1": f1}

# 学習パラメータの設定
batch_size = 16
logging_steps = len(dataset_encoded["train"]) // batch_size
model_name = "sample-text-classification-bert"

training_args = TrainingArguments(
    output_dir=model_name,
    num_train_epochs=2,
    learning_rate=2e-5,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    weight_decay=0.01,
    evaluation_strategy="epoch",
    disable_tqdm=False,
    logging_steps=logging_steps,
    push_to_hub=False,
    log_level="error",
)

# Trainerの定義
trainer = Trainer(
    model=model,
    args=training_args,
    compute_metrics=compute_metrics,
    train_dataset=dataset_encoded["train"],
    eval_dataset=dataset_encoded["validation"],
    tokenizer=tokenizer
)

# トレーニング実行
trainer.train()

この内容についての解説は以下の記事を参照ください。

TrainingArgumentの変更

EarlyStoppingを使用するには、基本的にはコールバックを設定することになりますですが、使用するためにはTrainingArgumentが特定の条件を満たすことが必要です。

今回は、以下の２点をTrainingArgumentで変更しました。

• load_best_model_at_end=Trueとする必要があったため指定
• evaluation_strategyとsave_strategyが同じである必要があったため指定
  • 正確には、これはload_best_model_at_end=Trueとする際の条件

またnum_train_epochs=2ではEarlyStoppingが働く前に学習が終わりますので、num_train_epochs=20に変更しています。

training_args = TrainingArguments(
    output_dir=model_name,
    num_train_epochs=20,
    learning_rate=2e-5,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    weight_decay=0.01,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    disable_tqdm=False,
    logging_steps=logging_steps,
    push_to_hub=False,
    log_level="error",
    load_best_model_at_end=True,
)

EarlyStoppingコールバックを指定

EarlyStoppingのコールバックをTrainerクラスのcallbacksに指定します。

early_stopping_patienceは、どの程度eval_datasetで計算される指標が改善しなかったら学習を停止するかの設定で、今回は3(この場合単位はエポックとなる)としています。

from transformers import EarlyStoppingCallback

trainer = Trainer(
    model=model,
    args=training_args,
    compute_metrics=compute_metrics,
    train_dataset=dataset_encoded["train"],
    eval_dataset=dataset_encoded["validation"],
    tokenizer=tokenizer,
    callbacks=[EarlyStoppingCallback(early_stopping_patience=3)],
)

trainer.train()

なお、ベストと判断するための指標は、デフォルトではlossが使用されます。

別の指標（例えばF1スコアなど）を判断基準にしたい場合は、metric_for_best_model="f1"などと指定すればOKです。 この場合、本記事の例のようにcompute_metricsを同じ名前("f1")をreturnするように定義して与える必要があります。

まとめ

いかがでしたでしょうか？

今回は実際に業務で学習を行う際には必要となってくるEarly StoppingについてHugging Faceで実施する方法を紹介しました。 単純なコールバックのみの設定ではエラーとなり少し調査が必要でしたので今回記事にさせていただきました。

本記事がHugging Faceを使われる方の参考になれば幸いです。