Make組ブログ

Python、Webアプリや製品・サービス開発についてhirokikyが書きます。

HuggingFace Trainerで学習とハイパーパラメーターチューニングを両立するプログラムを書く

HuggingFace Trainer を使っていて、トレーニング用のプログラムとハイパーパラメーターチューニング用のプログラムを両立させる方法を考えてみました。

HuggingFace Trainerでは trainer.hyperparameter_search(...) メソッドを使ってハイパーパラメーターのチューニングが可能です。たとえばOptunaを使う場合も backend="optuna" と指定するだけで済み、Pruneの設定なども簡単に行えます。

レーニングもハイパーパラメーターチューニングも Trainer(...) を使いますので、ある程度処理を共通化したいところです。ですがいくつか違う点があるのでその差分をうまく吸収するコードを書いていきましょう。

たとえばこんな点があります:

  • .hyperparameter_search() 時はモデルを model= で渡すのではなく、 model_init= に関数を渡す
  • ハイパーパラメーターチューニング時に保存やevalの実行頻度を変える
    • save_strategy="no" を指定してモデルを保存しない(学習時は save_strategy="epoch"
    • eval_strategy="epoch" を指定(学習時は eval_strategy="steps"

レーニング用とチューニング用で再利用する

以下のように、 Trainer を作る関数を用意します。この関数の結果に trainer.train_model() するか trainer.hyperparameter_search() を呼ぶことになります。

def model_init():
    return AutoModel.from_pretrained(...)


def trainer_builder(
    epochs,
    output_dir,
    learning_rate,
    model_init=False,
    **trainer_arguments,
):
    # データセットの用意
    dataset = load_dataset(...)

    train_data = ...
    val_data = ...

    if model_init:
        model_kwargs = {"model_init": model_init, "model": False}
    else:
        model_kwargs = {"model": model_init()}

    args = TrainerArguments(
        output_dir=output_dir,
        num_train_epochs=epochs,
        learning_rate=learning_rate,
        per_device_train_batch_size=batch_size,
        per_device_eval_batch_size=batch_size,
        **trainer_kwargs,
    )

    return Trainer(
        **model_kwargs,
        args=args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
    )

レーニング用のプログラムでは以下のように trainer_builder() を呼びだします。
コマンドの引数をうまく取ったり、そのまま受け渡したりしてハイパーパラメーターチューニング時との差分を埋めています。

@click.command()
@click.option("--epoch", type=int, default=32)
@click.option("--output_dir", default="output")
@click.option("--learning_rate", type=float, default=1e-05)
@click.option("--resume_ckpt", default=False, is_flag=True)
def main(resume_ckpt, **cmd_kwargs):
    trainer = trainer_builder(
        save_strategy="epoch",
        eval_strategy="steps",
        eval_steps=1000,
        **cmd_kwargs,
    )
    trainer.train(resume_from_checkpoint=resume_ckpt)

他にもコマンドの引数として受け取りたいパラメーターがあれば随時追加すればOKです。

ハイパーパラメーターチューニングは以下のようにします。

def hp(trial):
    return {
        "learning_rate": trial.suggest_float("learning_rate", 1e-05, 1e-04),
    }


def main():
    trainer = trainer_builder(
        epochs=10,
        output_dir="output",
        learning_rate=1e-05,  # どのみち上書きされる
        batch_size=32,
        model_init=True,
        save_strategy="no",
        eval_strategy="epoch",
        report_to=[],
    )
    result = trainer.hyperparameter_search(
        hp,
        direction="minimize",
        n_traials=10,
        backend="optuna",
        pruner=optuna.pruners.MedianPruner(),
    )
    pprint(result)

これでトレーニング用とチューニング用でコードをある程度共通化できます。

注釈

この書き方が万全という自信はなく、まだ実験中なので他に良い書き方があれば教えてください。たとえばTrainerを共通化するのではなく、データセットの準備を共通化すれば済むなどの考え方は他にもありそうです。また、トレーニング用のコードでは model_init の代わりに model 引数を使っていますが、 前者で済むのであれば切り替えは不要です(.hyperparameter_search() の場合は model_init が必須です)。

おわりに

Shodoではモデルのファインチューニングを行っており、自社でAIを運用しております。モデルのハイパーパラメーターチューニングや、継続的な開発をしやすいプログラムが重要になりますので、こういった工夫が重要になります。ぜひほかにも「こうしたほうが良いよ!」等のアドバイスがあればぜひ教えていただけると嬉しいです。

また現在、Shodoではアドベントカレンダー応援クーポンを配布しております。80%オフでShodoを最長3ヶ月間使えるクーポンです。以下のクーポンコードをご購入時に入力して、このアドベントカレンダーの季節にShodoのAI校正をブログの執筆にお役立てください。

XMAS2024

shodo.ink

執筆:@hirokiky
Shodoで執筆されました