Search Space Basics#

  • This chapter covers the implementation of search space for hyperparameter optimization

The search space refers to the range of values that can be used to explore each hyperparameter during HPO.

SearchSpace Class#

The SearchSpace class in the Ablator library is responsible for defining the search space for hyperparameters. It allows you to specify the range for continuous hyperparameters, the categorical values for discrete hyperparameters, and the data type of the hyperparameter.

Import SearchSpace using from ablator.main.configs import SearchSpace.

The SearchSpace class takes the following arguments:

  • value_range: This argument is used to define the range for continuous hyperparameters. It is specified in the format of [lower_bound, upper_bound]. For example, you may set value_range = [0, 1.0] for the dropout layer.

  • categorical_values: This argument is used for discrete hyperparameters. For example, to test the model’s performance on different batch sizes, we can use categories like [32, 64, 128], etc.

  • value_type: This argument defines the data type of the hyperparameter. There are two data types supported: “int” for integer values and “float” for decimal or floating-point values. For example, value_type = “int” for integer type.

Note that categorical values do not require a value type.

from ablator.main.configs import SearchSpace

SearchSpace(value_range=[0.05, 0.1], value_type="float")
SearchSpace(categorical_values=[32, 64, 128])

Creating a search space for hyperparameters.#

The goal of defining a search space in hyperparameter optimization (HPO) is to encapsulate the possible values and ranges of hyperparameters.

We use a Python dictionary where the key represents a hyperparameter in the config object, and the value is the SearchSpace object.

search_space_lr = {
    "train_config.optimizer_config.arguments.lr": SearchSpace(
        value_range=[0.05, 0.1], value_type="float"
    )
}

Ablator can provide SearchSpace for:

  • Predefined Configurations in Ablator: Ablator offers predefined configurations for optimizers, schedulers, batch size, epochs, and more. These configurations are readily available for users to use in their experiments.

  • Custom Configurations added by users: Users can define custom configurations for parameters specific to their experiments. For example, parameters of custom models, activation functions, dropout layers, and other relevant hyperparameters.

Using “SearchSpace” for predefined configurations#

  • With optimizers

The ablator provides three predefined optimizers with their arguments. [SGD, Adam, and AdamW].

Note: You can find all the default values for each optimizer in the “Configuration Basics” chapter.

Specifically for AdamW, a search space may be:

search_space = {
    "train_config.optimizer_config.arguments.lr": SearchSpace(
        value_range  = [0.01, 0.05],
        value_type = "float"
    ),
    "train_config.optimizer_config.arguments.eps": SearchSpace(
        value_range  = [1e-9, 1e-7],
        value_type = "float"
    ),
    "train_config.optimizer_config.arguments.weight_decay": SearchSpace(
        value_range  = [1e-4, 1e-3],
        value_type = "float"
    ),
}

  • With schedulers

The ablator also provides three schedulers: [step, cycle, and plateau]. You can use the SearchSpace with schedulers for their respective arguments.

For example, a search_space for scheduler “plateau” will look like this:

search_space = {
    "train_config.scheduler_config.arguments.min_lr": SearchSpace(
        value_range  = [1e-6, 1e-4],
        value_type = "float"
    ),
    "train_config.scheduler_config.arguments.mode": SearchSpace(
        categorical_values = ["min", "max", "auto"]
    ),
    "train_config.scheduler_config.arguments.threshold": SearchSpace(
        value_range  = [1e-5, 1e-3],
        value_type = "float"
    ),
}

  • Other parameters

We can also provide SearchSpace to other parameters like epochs, batch_size, etc. inside train_config.

For example, trying different batch_size or epochs can be easily done using the code:

search_space = {
    "train_config.batch_size": SearchSpace(
        categorical_values = [32, 64, 128]
    ),
    "train_config.epochs": SearchSpace(
        value_range  = [10, 20],
        value_type = "int"
    ),
}

Using SearchSpace for Custom Configs#

In this example, we will explore the “hidden_size” and “activation” parameters of a custom model and will create a search space for it.

Suppose, a custom model config takes some inputs like:

class CustomModelConfig(ModelConfig):
    ...
    input_size :int
    hidden_size :int
    num_classes :int
    activation: str
    ...

model_config = CustomModelConfig(
    input_size = 28*28,
    hidden_size = 256,
    num_classes = 10,
    activation = "relu"
)

Remember, a model config is defined to pass its arguments inside the constructor of the PyTorch model. For example:

class MyModel(nn.Module):
    def __init__(self, config: CustomModelConfig) -> None:
        activation_list = {"relu" : nn.ReLU(), "elu": nn.ELU()}

        self.fc1 = nn.Linear(config.input_size, config.hidden_size)
        self.act1 = activation_list[config.activation]
        ...

We can provide a SearchSpace for model’s “hidden_size” parameter like this:

search_space = {
    "model_config.hidden_size": SearchSpace(
        value_range=[250, 500], value_type="int"
    ),
}

Similarly, we can also provide a search space for “activation”.

search_space = {
    "model_config.activation": SearchSpace(
        categorical_values = ["relu","elu"]
    ),
}

The ablator will create trials using different “hidden_size” or “activations” according to the “search_space” provided.

Finally, the search_space dictionary is passed to the ParallelConfig, which will be explored in detail in the HPO tutorial.

SearchSpace for YAML files#

If we are using a YAML file to define configurations, we can specify a search space as follows:

...
search_space:
  model_config.hidden_size:
    value_range:
    - '250'
    - '500'
    categorical_values: null
    value_type: int
  train_config.optimizer_config.arguments.lr:
    value_range:
    - '0.001'
    - '0.01'
    categorical_values: null
    value_type: float
  model_config.activation:
    value_range: null
    categorical_values:
    - relu
    - leakyRelu
    - elu
    value_type: float
...

Conclusion#

We have successfully explored the SearchSpace class and various ways to utilize it. In the subsequent chapter, we will learn how to use search_space with “ParallelConfig” for HPO.