theislab / cpa

The Compositional Perturbation Autoencoder (CPA) is a deep generative framework to learn effects of perturbations at the single-cell level. CPA performs OOD predictions of unseen combinations of drugs, learns interpretable embeddings, estimates dose-response curves, and provides uncertainty estimates.
BSD 3-Clause "New" or "Revised" License
83 stars 17 forks source link

Tuner isn't logging plan_kwargs on WandB #47

Closed ArianAmani closed 6 months ago

ArianAmani commented 6 months ago

In the hyperparameter tuning code, here: https://github.com/theislab/cpa/blob/f415f9f29535b1da153de3c80aa16689d4a2d400/cpa/_tuner.py#L568 when we pop the plan arguments from train_args, it results in plan arguments being ignored while reporting the parameters and you won't find their values in the wandb report analysis.

Possible fix: Replace https://github.com/theislab/cpa/blob/f415f9f29535b1da153de3c80aa16689d4a2d400/cpa/_tuner.py#L567-L574 with

actual_train_args = {}
for key in train_args.keys():
    if key in plan_kwargs_keys:
        plan_kwargs[key] = train_args[key]
    else:
        actual_train_args[key] = train_args[key]

train_args = {
    "enable_progress_bar": True,
    "logger": experiment.get_logger(get_context().get_trial_name()),
    "callbacks": [experiment.metrics_callback],
    **actual_train_args,
}