Issue with Optuna visualization in web browser

🤗Transformers
1

Hello everyone,

I am using the Optuna library for hyperparameter optimization in a language model training script. I am trying to visualize the optimization results using Optuna’s plot_optimization_history function and then saving the figure as an HTML file using fig.write_html("optuna_visualization.html").

The HTML file is successfully created and opens correctly in all browsers. However, the figure does not show any data. In the terminal, the best hyperparameters are printed correctly.

Here is a snippet of my code:

        # Ask the user if they would like to use Optuna for hyperparameter optimization
        optuna_mode = input("\n➤Would you like to use the default or Optuna settings to train? (\033[94mpadrao\033[0m/\033[91moptuna\033[0m): ")
        print()
  
        study_name = "tess"
        db_path = "example.db"
        if optuna_mode.lower() == 'optuna':
            study = optuna.create_study(study_name=study_name, storage=f"sqlite:///{db_path}", direction="minimize")
            study.optimize(lambda trial: objective(trial, model, tokenized_datasets, data_collator, tokenizer), n_trials=2)

            if len([t for t in study.trials if t.state == optuna.trial.TrialState.COMPLETE]) > 0:
                best_params = study.best_params
                print("\nBest hyperparameters: ", best_params)
                # View the results
                fig = vis.plot_optimization_history(study)
                fig.write_html("optuna_visualization.html")
                fig.show()
            else:
                print("No attempts were completed.")

And the objective function:

# Function for hyperparameter optimization with Optuna
def objective(trial, model, tokenized_datasets, data_collator, tokenizer):
    learning_rate = trial.suggest_float("learning_rate", 1e-6, 1e-3, log=True)
    num_train_epochs = trial.suggest_int("num_train_epochs", 2, 4)
    weight_decay = trial.suggest_float("weight_decay", 1e-12, 1e-1, log=True)
    adam_epsilon = trial.suggest_float("adam_epsilon", 1e-10, 1e-6, log=True)
    per_device_train_batch_size = trial.suggest_categorical("per_device_train_batch_size", [8, 16, 32])

    print(f"\nHyperparameters Suggestion: learning_rate={learning_rate}, num_train_epochs={num_train_epochs}, weight_decay={weight_decay}, per_device_train_batch_size={per_device_train_batch_size}\n")

    training_args = TrainingArguments(
        output_dir=model_dir,
        overwrite_output_dir=True,
        num_train_epochs=num_train_epochs,
        per_device_train_batch_size=per_device_train_batch_size,
        learning_rate=learning_rate,
        weight_decay=weight_decay,
        adam_epsilon=adam_epsilon,
        save_steps=1_000,
        save_total_limit=2,
        prediction_loss_only=True,
        logging_steps=10,
        logging_dir='./logs',
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_datasets["train"],
        data_collator=data_collator,
        tokenizer=tokenizer,
    )

    #logging.getLogger("transformers.trainer").setLevel(logging.ERROR)

    # Train the model
    trainer.train()

    eval_loss = None
    if trainer.eval_dataset is not None:
        eval_loss = trainer.evaluate()["eval_loss"]
        print(f"Eval loss: {eval_loss}")

    return eval_loss if eval_loss is not None else float('inf')

I have checked that the study object contains the correct optimization results and that the visualization is being generated correctly. When I open the HTML file in the browser and check the developer console, I see no JavaScript errors.

Has anyone encountered a similar issue or have any idea how I can resolve this?

Thank you!