Hello all,
I am trying to perform hyperparameter tuning using optuna as a backend and log all trials to MLFlow. So I am initializing all parameters for mlflow and I also use trainer.hyperparameter_search(). My code is the follwing:
os.environ["MLFLOW_EXPERIMENT_NAME"] = "trainer-mlflow-demo"
os.environ["MLFLOW_FLATTEN_PARAMS"] = "1"
training_args = TrainingArguments(
output_dir=SAVE_DIR,
per_device_train_batch_size=PER_DEVICE_TRAIN_BATCH,
per_device_eval_batch_size=PER_DEVICE_EVAL_BATCH,
evaluation_strategy = "epoch",
logging_strategy="epoch",
save_strategy = "epoch") # need save strategy order to log the trials in optuna
trainer = Trainer(
model=None,
args=training_args,
train_dataset=tokenized_dset['train'],
eval_dataset=tokenized_dset['validation'],
model_init=model_init
)
best_trial = trainer.hyperparameter_search(
direction="minimize",
backend="optuna",
hp_space=optuna_hp_space,
n_trials=NUM_TRIALS
)
However, I face a problem when logging the parameters of each trial. I have observed that only the parameters of the 1st trial are logged and for the following trials, only the metrics are logged. There are other details missing in the following trials, such as the Run Name. In the image below, you can see the two trials I performed, with the one with full parameters being the 1st trial.
Has anyone seen this problem again? Is there a workaround to log all parameters in all trials?
Thank you in advance,
Petrina