Trainer: How to find the best learning rate?

Hi,

I have trained a text classification model with Huggingface Trainer and now want to do hyperparameter-tuning. As a first step I would like to find the best value for the “learning_rate” parameter. Is there maybe an automatic way to find the optimal learning rate?

Here’s a code snippet:

    def compute_metrics(p):
        pred, labels = p
        pred = np.argmax(pred, axis=1)

        accuracy = accuracy_score(y_true=labels, y_pred=pred)
        recall = recall_score(y_true=labels, y_pred=pred,average="macro")
        precision = precision_score(y_true=labels, y_pred=pred,average="macro")
        f1 = f1_score(y_true=labels, y_pred=pred,average="macro")

        return {"accuracy": accuracy, "precision": precision, "recall": recall, "f1": f1}

    args = TrainingArguments(
        # output_dir: directory where the model checkpoints will be saved.
        output_dir=f"{ansatz}_res",
        overwrite_output_dir=True,
        evaluation_strategy="steps",
        eval_steps=50,
        logging_strategy="steps",
        logging_steps=50,
        save_strategy="steps",
        #save_steps=200,
        learning_rate=5e-5,
        per_device_train_batch_size=8,
        per_device_eval_batch_size=8,
        num_train_epochs=3,
        load_best_model_at_end=True,
        metric_for_best_model="f1",
        # Paper hat LR von 6e-5 und 15 epochs: ComMa Hasoc 2020
    )

    trainer = Trainer(
        model=model,
        args=args,
        train_dataset=train_dataset,
        eval_dataset=val_dataset,
        compute_metrics=compute_metrics,
        callbacks=[EarlyStoppingCallback(early_stopping_patience=5)],
    )

Thanks in advance!

1 Like