Optunaのハイパラ探索時に、複数の指標を保存しておく方法

はじめに

Optunaでは、例えばAccuracyを最大化するようなハイパラ探索を行います。その場合、Accuracyをリアルタイムで表示したり、studyオブジェクトに保存したりすることができます。しかし、Accuracy以外にも、例えばPrecisionやRecallなどの指標を保存したい場合があります。そこで、今回はOptunaのハイパラ探索時に、複数の指標を保存しておく方法を紹介します。(注意点:複数指標を使ってハイパラ探索を行うという話ではありません。)

使用したバージョン

• python: 3.8.11
• optuna: 3.2.0

コード

通常の使い方

まずは、通常の使い方を紹介します。以下のコードは、ハイパラ探索の結果を見るために、全てのtrialが終わるまで待つコードです。

import optuna

def objective(trial):
    #学習データとテストデータを作成。make_data()は適当な関数.
    X_train, y_train, X_test, y_test = make_data()
    #あるtrialの学習とテストを行う関数
    accuracy = train_test(X_train, y_train, X_test, y_test)

    # Accuracyに基づいてハイパラ探索を行うように、accuracyを返す
    return accuracy

# 100通りのハイパラを探索する
study = optuna.create_study(direction='maximize')
study.optimize(objective, n_trials=100)
# studyを保存
with open('study.pkl', 'wb') as f:
    pickle.dump(study, f)

複数の指標を保存する使い方

trial毎に特定の指標を保存するには、set_user_attr()メソッドを使います。以下のコードは、各trialが終わるたびに、そのtrialのAccuracyとPrecisionとRecallを保存するコードです。

def objective(trial):
    #学習データとテストデータを作成。make_data()は適当な関数.
    X_train, y_train, X_test, y_test = make_data()
    #あるtrialの学習とテストを行う関数
    # 今回は、AccuracyとPrecisionとRecallを返すようにする
    accuracy,precision,recall = train_test(X_train, y_train, X_test, y_test)

    # trial毎にAccuracyとPrecisionとRecallを保存する
    trial.set_user_attr("precision", precision)
    trial.set_user_attr("recall", recall)
  
    # Accuracyに基づいてハイパラ探索を行うように、accuracyを返す
    return accuracy

# 100通りのハイパラを探索する
study = optuna.create_study(direction='maximize')
study.optimize(objective, n_trials=100)

# studyを保存
with open('study.pkl', 'wb') as f:
    pickle.dump(study, f)

まとめ

Optunaのハイパラ探索時に、複数の指標を保存しておく方法を紹介しました。 上記のようにtrialからset_user_attr()メソッドを叩くことで、trialと保存したい指標を紐づけることができます。修正点やご意見などありましたら、コメントお願いします。