How to get the best model from optuna hyper parameter search

Hi all,
I’ve been playing with optuna for hyperparmeter search. I’m using it with huggingface trainer and wandb.

I trained my models with the following codes:

best_run = trainer.hyperparameter_search(
    direction="maximize", 
    n_trials=4,
    hp_space=my_hp_space_optuna, 
    compute_objective=my_objective
)

What I want to do is to get the best model so that I can make predictions on my test data using it. Thus I’ve tried with the codes below but faced some issues :

best_run.predict(encoded_dataset['test'])
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-88-2f2e98dce207> in <module>()
----> 1 best_run.predict(encoded_dataset['test'])

AttributeError: 'BestRun' object has no attribute 'predict'

Also, I had a go with the following codes:

trainer.predict(encoded_dataset['test'])

It did run but it made predictions by using the latest trained model.

Any suggestion about this issue? Thanks in advance.

  • training arguments codes:
args = TrainingArguments(
    output_dir=output_dir,
    overwrite_output_dir=True,
    save_strategy = "epoch", 
    load_best_model_at_end=True,
    num_train_epochs=4,
    per_device_train_batch_size=64,
    per_device_eval_batch_size=32,
    learning_rate=2e-5, 
    metric_for_best_model= 'eval_pearsonr', 
    evaluation_strategy = "epoch",
    logging_dir=logging_dir,
    logging_strategy="steps",
    logging_steps=1,
    report_to="wandb",
)