Using hyperparameter-search in Trainer

🤗Transformers
55

Hi all!

I’m using Optuna for hyper-parameter search, but I have a doubt/problem.

Here the performance report obtained during fine-tuning a pretrained BERT for Polarity Classification:

Epoch Training Loss Validation Loss Accuracy F1 Precision Recall
1 0.546600 0.438429 0.795547 0.749214 0.745114 0.753995
2 0.277000 0.457287 0.786100 0.742984 0.735595 0.753490
3 0.207700 0.595937 0.800270 0.743335 0.751293 0.736927

Optuna says:

[I 2021-07-23 10:27:45,517] Trial 0 finished with value: 0.7433345815390376 and parameters: {‘learning_rate’: 4.5470513013108546e-05, ‘warmup_steps’: 0.6, ‘weight_decay’: 0.07700782690380507}. Best is trial 0 with value: 0.7433345815390376.

So it seems that as best value for the run it takes the score of the last row, even tough I specified load_best_model=True and metric_for_best_model="eval_f1"

Here my code snippet:

from sklearn.metrics import precision_recall_fscore_support, accuracy_score

def compute_metrics(pred):
    labels = pred.label_ids
    preds = pred.predictions.argmax(-1)
    precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average='macro')
    acc = accuracy_score(labels, preds)
    return {
        'accuracy': acc,
        'f1': f1,
        'precision': precision,
        'recall': recall
    }

class MemorySaverCallback(TrainerCallback):
    "A callback that deleted the folder in which checkpoints are saved, to save memory"
    def __init__(self, run_name):
        super(MemorySaverCallback, self).__init__()
        self.run_name = run_name

    def on_train_begin(self, args, state, control, **kwargs):
        print("Removing dirs...")
        if os.path.isdir(f'./{self.run_name}'):
            import shutil
            shutil.rmtree(f'./{self.run_name}')
        else:
            print("\n\nDirectory does not exists")

training_args = TrainingArguments(
    RUN_NAME, 
    num_train_epochs=15,
    per_device_train_batch_size=64,
    per_device_eval_batch_size=64,
    evaluation_strategy="epoch",
    logging_strategy="steps",
    logging_steps=1,
    logging_first_step=False,
    overwrite_output_dir=True,
    save_strategy="no",
    save_total_limit=1,
    load_best_model_at_end=True,
    metric_for_best_model="eval_f1",
)

trainer = Trainer(
    model_init=partial(MyNet,2),
    args=training_args, 
    train_dataset=training_opos.select(range(2000)), 
    eval_dataset=validating_opos,
    compute_metrics=compute_metrics,
    callbacks=[EarlyStoppingCallback(early_stopping_patience=2), MemorySaverCallback(RUN_NAME)]
)

def my_hp_space_optuna(trial):
    return {
        "learning_rate": trial.suggest_float("learning_rate", 2e-6, 2e-4, log=True),
        "warmup_steps":  trial.suggest_float("warmup_steps", 0., 0.9, step=0.3),
        "weight_decay":  trial.suggest_float("weight_decay", 1e-6, 1e-1)
    }
def my_objective(metrics):
    return metrics["eval_f1"]

sa = trainer.hyperparameter_search(
    direction="maximize", 
    n_trials=1,
    hp_space=my_hp_space_optuna, 
    compute_objective=my_objective
)

Optuna version=2.8.0
Transformers version=4.6.1

Thanks in advance!

3 Likes