How to integrate TorchTrainer with Ray Tune and HyperOpt space?

TorchTrainer expects the tuning configuration to be provided in the {"train_loop_config": config} format as shown in this example.

However, when using custom search spaces, e.g., one that is defined with hyperopt, the config is automatically generated by Ray and passed to the trainer function, so it does not conform to the {"train_loop_config": config} format.

A quick and dirty solution to this problem is adding something like self.config = {"train_loop_config": self.config} before the dictionaries are merged here.

Is there a cleaner way to pass configurations generated from spaces to TorchTrainer?

Do param_space={"train_loop_config": {"a": tune.grid_search([1, 2])}} and Tuner(..., searcher_alg=HyperOptSearch()) work for you?

1 Like

It will probably work. I was just curious to see if it could be done automatically without changing the param_space.

The way I currently use custom hyperopt spaces does not work with Tune out of the box. For example, assume the following conditional search space:

space = {
    ##### Augmentation ######
    # Cutmix
    "cutmix": hp.choice(
        "ctmx",
        [
            {
                "cutmix": True,
                "cutmix_beta": 1.0,
                "cutmix_prob": hp.quniform("cutmix_prob", 0.10, 0.50, 0.05),
            },
            {"cutmix": False},
        ],
    ),
}

The configuration that Tune passes to workers looks like this:

{"cutmix": {"cutmix": False}}

If your code is designed to receive cutmix, cutmix_beta and cutmix_prob, the provided config will not work, i.e., you need to process spaces defined in Tune, which only have the expected keys, and the ones defined with hyperopt, which may have nested dictionaries, differently.

So I have written a helper function to first flatten the nested dictionaries, and use the updated config in the worker.

def hyperopt_to_ray(config):
    # Flatten nested config when HyperOpt custom space is defined
    outer_keys = []
    for key in config:
        if isinstance(config[key], dict):
            outer_keys.append(key)
    for key in outer_keys:
        inner_dict = config.pop(key)
        config.update(inner_dict)

Changing the param_space definition to something like param_space = {"train_loop_config": param_space} would make this whole process of dealing with spaces more complicated.