ashleve / lightning-hydra-template

PyTorch Lightning + Hydra. A very user-friendly template for ML experimentation. ⚡🔥⚡
4.27k stars 654 forks source link

How to define a conditional search space ? #457

Open tianshuocong opened 2 years ago

tianshuocong commented 2 years ago

Hi!

I want to search the best optimizer for the given "mnist_example" from SGD and Adam. However, for SGD, I also want to know which momentum value is the best (which Adam doesn't need), but for Adam, I need search beta.

Therefore, how can we define such a hparams_search/name.yaml file?

hydra:
  ...
  sweeper:
  ...
    params:
      model.optimizer._target_: choice("torch.optim.SGD", "torch.optim.Adam")
       ......

Thanks so much!

ashleve commented 2 years ago

None of the available in hydra optimizers support conditional search space as far as I'm aware.

You should either write custom optimization pipeline using framework which supports it, or just sample both hyperparameters at once and choose the right parameter in code when instantiating optimizer.

tianshuocong commented 2 years ago

Thanks for your reply!

However in optuna, RandomSampler and TPESample support conditional search space, so I am trying to find a solution with your template.

https://optuna.readthedocs.io/en/stable/reference/samplers/index.html

ashleve commented 2 years ago

My bad, thanks for letting me know!

I think this feature was implemented in this PR: https://github.com/facebookresearch/hydra/pull/1909

So you can extend the search space with a custom python function which should allow for conditions.

Example: https://github.com/facebookresearch/hydra/blob/5300a8632d9e5816863b37dbcf32a3fc5326a3c6/plugins/hydra_optuna_sweeper/example/custom-search-space/config.yaml#L19 https://github.com/facebookresearch/hydra/blob/5300a8632d9e5816863b37dbcf32a3fc5326a3c6/plugins/hydra_optuna_sweeper/example/custom-search-space-objective.py#L16-L23

tianshuocong commented 2 years ago

Thank for your examples!

Here I am also confused how to set momentum value for the torch.optim.SGD optimizer?

ashleve commented 2 years ago

I don't yet have experience with conditional search space in Optuna but from what I understand, you can use the pythonic search space in configure(cfg, trial) method as in official hydra examples, so I assume you should be able to do what's explained in Optuna documentation here: https://optuna.readthedocs.io/en/stable/tutorial/10_key_features/002_configurations.html

Example of conditional search space from docs:

def objective(trial):
    classifier_name = trial.suggest_categorical("classifier", ["SVC", "RandomForest"])
    if classifier_name == "SVC":
        svc_c = trial.suggest_float("svc_c", 1e-10, 1e10, log=True)
        classifier_obj = sklearn.svm.SVC(C=svc_c)
    else:
        rf_max_depth = trial.suggest_int("rf_max_depth", 2, 32, log=True)
        classifier_obj = sklearn.ensemble.RandomForestClassifier(max_depth=rf_max_depth)
ashleve commented 2 years ago

It's probably worth mentioning this somewhere in the template readme.

tianshuocong commented 2 years ago

Hi!

Really thanks for your reply!

However, I am so sorry that I still cannot solve this problem.

Let me rephrase my question: I want to use the hydra-optuna to search which optimizer is better from [SGD, Adam]. Meanwhile, when the selected optimizer is SGD, I also want to search the best momentum (which Adam do not need). Therefore, I would like to construct a conditional search space.

Here is an example, where line 81-83 show that

  optimizer_name = trial.suggest_categorical("optimizer", ["Adam", "RMSprop", "SGD"])
  lr = trial.suggest_float("lr", 1e-5, 1e-1)
  optimizer = getattr(optim, optimizer_name)(model.parameters(), lr=lr)

So, I think I can add such codes in src/models/mnist_module.py to define the optimizer like

      optim_name = self.hparams.optimizer.func.__name__
      lr = Trial.suggest_float("lr", 1e-5, 1e-1)
      mo = Trail.suggest_float("mo", 0.1, 0.9)

      if optim_name == 'Adam':
          optimizer = self.hparams.optimizer(params=self.parameters(), lr = lr)
      elif optim_name == 'SGD':
          optimizer = self.hparams.optimizer(params=self.parameters(), lr = lr, momentum=mo)

However, I do not know how to use trail.

here said "Hydra's Optuna Sweeper allows users to provide a hook for custom search space configuration. This means you can work directly with the optuna.trial.Trial object to suggest parameters. "

I try from optuna.trial import Trial, it shows suggest_float() missing 1 required positional argument: 'high'.

So can you show me a detail example? Really thanks for your help!

ashleve commented 2 years ago

@tianshuocong from what I understand, you're not sure how to access the Trial object in your code?

  1. Create a def configure(cfg, trial) method somewhere in your code. As you can see this method has trial object among the args. Example: https://github.com/facebookresearch/hydra/blob/5300a8632d9e5816863b37dbcf32a3fc5326a3c6/plugins/hydra_optuna_sweeper/example/custom-search-space-objective.py#L16-L23

  2. Add custom_search_space=... to your hydra sweeper config which tells hydra where is located the method for configuring the trial, e.g. custom_search_space=src.train.configure. Example: https://github.com/facebookresearch/hydra/blob/5300a8632d9e5816863b37dbcf32a3fc5326a3c6/plugins/hydra_optuna_sweeper/example/custom-search-space/config.yaml#L19

elidub commented 2 weeks ago

Thank you for working on this feature! I have looked at #1909 and it's documentation, but I can't make it work with the dropout usecase as illustrated in the first message of this thread.

Question

It would be nice if one could setup all the ranges at one central place (the Hydra config), and set the logic in the custom_search_space.

Consider the minimal working example:

config.yaml

defaults:
  - override hydra/sweeper: optuna

hydra:
  sweeper:
    sampler:
      seed: 123
    direction: minimize
    study_name: custom-search-space
    storage: null
    n_trials: 20
    n_jobs: 1

    params:
      use_dropout: True, False
      dropout: interval(0.0, 0.5) # dropout should not be considered as a hparam when use_dropout is False
    custom_search_space: run.configure

use_dropout: False
dropout: 0.

run.py

import hydra
from omegaconf import DictConfig
from optuna.trial import Trial

@hydra.main(config_path="", config_name="config", version_base=None)
def main(cfg: DictConfig) -> float:
    print(f"{cfg =          }")

    loss: float = 1.

    if cfg.use_dropout:
        # simulate that using dropout DECREASES the performance
        loss += cfg.dropout * 10.

    print('Done!')
    return loss

def configure(cfg: DictConfig, trial: Trial) -> None:

    use_dropout: bool = cfg.use_dropout

    if not use_dropout:
        # How do I remove (or set to zero) the dropout parameter?
        trial.suggest_float("dropout", 0.0, 0.0) # No effect? 

        # Other things I've tried:
        # trial.params['dropout'] = 0.0 # No effect?
        # trial.suggest_categorical("dropout", [0.0]) # Can't change distribution.

    print(f"{trial.params = }")

if __name__ == '__main__':
    main()

The print statements are merely for debugging.

How do I disable dropout from the trial if use_dropout = False?

If the above is not possible, then an alternative might the the following.

Alternative question

It might be possible to set the interval in configure() instead of in config.yaml. However, that also doesn't seem to work.

config.yaml the same as above.

run.py

import hydra
from omegaconf import DictConfig
from optuna.trial import Trial

@hydra.main(config_path="", config_name="config", version_base=None)
def main(cfg: DictConfig) -> float:
    print(f"{cfg =          }")

    loss: float = 100.

    if cfg.use_dropout:
        # simulate that using dropout INCREASES the performance
        loss -= cfg.dropout * 10.

    print('Done!')
    return loss

def configure(cfg: DictConfig, trial: Trial) -> None:

    use_dropout: bool = cfg.use_dropout

    if use_dropout:
        # Is not working. The dropout parameter from config.yaml is not being replaced.
        trial.suggest_float("dropout", 1.0, 5.0)

    print(f"{trial.params = }")

if __name__ == '__main__':
    main()

Note that in this example the dropout increases the performance, while in the previous question it was decreasing.

However, the dropout range from config.yaml is not being replaced when use_dropout = True. Even when I remove/comment the dropout: interval(0.0, 0.5) line from config.yaml, it stays the default value of dropout: 0..

If I remove dropout: interval(0.0, 0.5) from config.yaml and remove the if statement if use_dropout, the range is being replaced by trial.suggest_float("dropout", 1.0, 5.0), but that would remove the necessary logic...

Any ideas or suggestions how to solve this?

MattiasDC commented 2 weeks ago

Hi @elidub,

I wrote that custom search space logic a few years ago and got a notification on this question. Note that at the point the custom search space is invoked, you should retrieve the already configured parameters from the trial object, and not from the config. In the config.yaml you've set in these two lines:

use_dropout: False
dropout: 0.

And the ConfigDict in configure, will return these values when you ask them. Instead, you should ask the trial, as that is what is actually going to be used for the run. Also note that you should not define dropout in both the search space in the yaml and in your custom search space, as then the yaml will take precedence.

This should do what you want, or at least get you started.

defaults:
  - override hydra/sweeper: optuna

hydra:
  sweeper:
    sampler:
      seed: 123
    direction: minimize
    study_name: custom-search-space
    storage: null
    n_trials: 20
    n_jobs: 1

    params:
      use_dropout: True, False
    custom_search_space: run.configure

use_dropout: True
dropout: 0.
import hydra
from omegaconf import DictConfig
from optuna.trial import Trial

@hydra.main(config_path="", config_name="config", version_base=None)
def main(cfg: DictConfig) -> float:
    print(f"{cfg =          }")

    loss: float = 100.

    if cfg.use_dropout:
        # simulate that using dropout INCREASES the performance
        loss -= cfg.dropout * 10.

    print('Done!')
    return loss

def configure(cfg: DictConfig, trial: Trial) -> None:
    use_dropout: bool = trial.params['use_dropout']

    if use_dropout:
        # Is not working. The dropout parameter from config.yaml is not being replaced.
        trial.suggest_float("dropout", 1.0, 5.0)
    else:
        # If you need dropout to be always defined for further processing, you can remove the remark below
        #trial.set_user_attr("dropout", 0.)
        pass

    print(f"{trial.params=}")

if __name__ == '__main__':
    main()

You can now run this with: python run.py --multirun

elidub commented 2 weeks ago

@MattiasDC Thank you for your reply.

I think trial.set_user_attr("dropout", 0.) does not work as expected. I'm not even sure if it does anything. If we set dropout: 1. in config.yaml and have the if-else statement as you suggested:

    if use_dropout:
        trial.suggest_float("dropout", 1.0, 5.0)
    else:
        trial.set_user_attr("dropout", 0.)

The dropout is actually still 1. instead of 0. if use_dropout = False, seeing from the print statement. When having trial.suggest_float("dropout", 0., 0.) instead of trial.set_user_attr("dropout", 0.) it seems to work as intended.

MattiasDC commented 2 weeks ago

Hi @elidub,

In that case I would refrain from using suggest_float with no actual range (you don't need dropout parameter when not using it), as this will make it more difficult for finding an optimum for Optuna.