Search space for different types of optimizers and schedulers#

Different optimizers have different update rules and behavior, and they may perform better or worse depending on the specific dataset and model architecture. Hence, trying out different optimizers and learning rate schedulers can be beneficial for HPO.

  • To work with different optimizers effectively in the ablator, it is necessary to create custom OptimizerConfig objects that can handle passing either torch-defined or custom optimizers to the ablator.

  • This is similar to the schedulers.

Let us first import the necessary libraries and modules:

import torch
import torch.nn as nn
from torch.optim.lr_scheduler import OneCycleLR, ReduceLROnPlateau, StepLR

from ablator import configclass, ConfigBase

Search space for different optimizers#

We define a function called create_optimizer that creates an optimizer object based on the given inputs (optimizer name, model to optimize, and learning rate). In this example, we support three optimizers: Adam, AdamW, and SGD (but we can also pass our custom-defined optimizers). In specific, the function does the following:

  • Creates a list of model parameters parameter_groups from the model module model.named_parameters().

  • Defines dictionaries with specific parameters for each optimizer.

  • Create the optimizer using the model parameters, learning rate, and the defined dictionaries for each optimizer parameters.

Returns the optimizer object.

def create_optimizer(optimizer_name: str, model: nn.Module, lr: float):

    parameter_groups = [v for k, v in model.named_parameters()]

    adamw_parameters = {
      "betas": (0.0, 0.1),
      "eps": 0.001,
      "weight_decay": 0.1
    }
    adam_parameters = {
      "betas" : (0.0, 0.1),
      "weight_decay": 0.0
    }
    sgd_parameters = {
      "momentum": 0.9,
      "weight_decay": 0.1
    }

    Optimizer = None

    if optimizer_name == "adam":
        Optimizer = optim.Adam(parameter_groups, lr = lr, **adam_parameters)
    elif optimizer_name == "adamw":
        Optimizer = optim.AdamW(parameter_groups, lr = lr, **adamw_parameters)
    elif optimizer_name == "sgd":
        Optimizer = optim.SGD(parameter_groups, lr = lr, **sgd_parameters)


    return Optimizer

Finally, we create an Optimizer configuration CustomOptimizerConfig. Internally, Ablator requires that the optimizer config has function make_optimizer with input as a model module:

@configclass
class CustomOptimizerConfig(ConfigBase):
    name: Literal["adam", "adamw", "sgd"] = "adam"
    lr: float = 0.001

    def make_optimizer(self, model: nn.Module):
        return create_optimizer(self.name, model, self.lr)

optimizer_config = CustomOptimizerConfig(name = "adam", lr = 0.001)

  • Here the configuration attribute name will be used in the search space, and we’re allowing search space to be from the set of values ["adam", "adamw", "sgd"].

  • Inside make_optimizer, we call create_optimizer with the model, the name and lr attributes of the config object, and this function will return the corresponding optimizer.

Search space for different schedulers#

We define a function called create_scheduler that creates a scheduler object based on the given inputs (scheduler name, the model to optimize, the optimizer used). In this example, we support three schedulers: StepLR, OneCycleLR, and ReduceLROnPlateau (but we can also pass our custom-defined schedulers). In specific, the function does the following:

  • Defines dictionaries with specific parameters for each scheduler.

  • Create the scheduler using the optimizer and the defined dictionaries for each scheduler parameters.

  • Return the scheduler object.

We also define a second function called scheduler_arguments that returns the arguments of the scheduler

def create_scheduler(scheduler_name: str, model: nn.Module, optimizer: torch.optim):

  parameters = scheduler_arguments(scheduler_name)

  Scheduler = None

  if scheduler_name == "step":
    Scheduler = StepLR(optimizer, **parameters)
  elif scheduler_name == "cycle":
    Scheduler = OneCycleLR(optimizer, **parameters)
  elif scheduler_name == "plateau":
    Scheduler = ReduceLROnPlateau(optimizer, **parameters)

  return Scheduler

def scheduler_arguments(scheduler_name):
  if scheduler_name == "step":
    return {
      "step_size" : 1,
      "gamma" : 0.99
    }
  elif scheduler_name == "cycle":
    return {
      "patience":  10,
      "min_lr":  1e-5,
      "mode":  "min",
      "factor":   0.0,
      "threshold":  1e-4
    }
  elif scheduler_name == "plateau":
    return {
      "max_lr": 1e-3,
      "total_steps": 10
    }

Similarly, we also create a custom config CustomSchedulerConfig, defining the required method make_scheduler with shceduler name, the model, and the optimizer as inputs.

@configclass
class CustomSchedulerConfig(SchedulerConfig):
    name: Literal["step", "cycle", "plateau"] = "step"

    def __init__(self, name, arguments=None):
        self.arguments = scheduler_arguments(self.name)
        super(CustomSchedulerConfig, self).__init__(name=self.name, arguments=self.arguments)

    def make_scheduler(self, model: torch.nn.Module, optimizer: torch.optim):
        return create_scheduler(self.name, model, optimizer)

scheduler_config = CustomSchedulerConfig(name = "step")

  • Here the configuration attribute name will be used in the search space, and we’re allowing search space to be from the set of values ["step", "cycle", "plateau"].

  • We overwrite the constructor, creating an attribute called arguments, which is internally accessed by ablator and pass it to the parent class constructor.

  • Inside make_scheduler, we call create_scheduler with the optimizer, the name and lr attributes of the config object, and this function will return the corresponding scheduler.

Note

Remember to redefine the TrainConfig config class, hence the ParallelConfig, before creating the training config to pass in the optimizer and scheduler config objects. E.g:

@configclass
class CustomTrainConfig(TrainConfig):
  optimizer_config: CustomOptimizerConfig
  scheduler_config: CustomSchedulerConfig

@configclass
class CustomParallelConfig(ParallelConfig):
  model_config: CustomModelConfig
  train_config: CustomTrainConfig

Create search space for optimizers and schedulers#

Now, we can try out different optimizers and schedulers by providing a search space to the ablator.

search_space = {
    "train_config.optimizer_config.lr": SearchSpace(value_range = [0.001, 0.01], value_type = 'float'),
    "train_config.optimizer_config.name": SearchSpace(categorical_values = ["adam", "sgd", "adamw"]),
    "train_config.scheduler_config.name": SearchSpace(categorical_values = ["step", "cycle", "plateau"])
}

Note:

In the default optimizer config, providing the name of the optimizer in the config will create an object of the associated optimizer class. Simply changing the name in the search space will result in a mismatch in the class type, causing an error. Hence, we have to define custom configs in this way.

One benefit this method offers is that we can define our custom optimizers or schedulers as a class and pass them to their respective configs for the ablator to manage training.

Conclusion#

Finally, with this, we can now test different optimizers and schedulers for our model.