[RFC] Customizable steps of `LightGBMTuner` and `LightGBMTunerCV`

Motivation

The steps of LightGBMTuner are fixed, and users cannot add the parameters to tune. I saw users’ demand to customize the steps.

For example,

• @Amyantis submitted a feature request to change the number of trials at each step in #1550, and
• @kdlin wants to tune additional parameters such as max_depth and max_bin in the Gitter conversation.

I created this issue to aggregate such user needs, and I’d like to discuss the feasible and user-friendly implementation of customizable steps.

Description

Currently, I do not have any ideas for API design, but let me share some existing attempts.

Approach 1 I tried to extend LightGBMTuner to add a step to tune max_depth. See https://colab.research.google.com/drive/1vBfktSi7qXc8T4X6x1ib3IeqIlqcOAOH?usp=sharing. It was a bit complicated because I needed to modify two classes, i.e., LightGBMTuner and _OptunaObjective. I don’t think this is a convenient way to customize LightGBM tuner’s step, and we may need major refactoring of LightGBMTuner and _OptunaObjective.

Approach 2 As mentioned in this Gitter comment, @jeffzi created a prototype of StepwiseTuner, which accepts a dictionary containing step information. It also enables us to create a stepwise tuner for other libraries such as XGBoost and CatBoost.

[jeffzi]

Some time ago I prototyped a StepwiseTuner class that accepts a list of named Steps in its init.

It’s essentially an abstraction of LightGBMTuner that needs to be specialized in a dedicated subclass. The shared logic of looping through the steps, along with keeping track of the global time and trial budgets across steps is implemented in StepwiseTuner.tune().

This approach is flexible. The steps and their order can be customized, the time and trial budget per step as well.

Here is a draft of the public interfaces:

from optuna import distributions, logging, pruners, samplers

class BaseStep(ABC):
    def __init__(
        self,
        sampler: samplers.BaseSampler = None,
        pruner: pruners.BasePruner = None,
        n_trials: int = None,
        timeout: int = None,
    ):
        self.sampler = sampler
        self.n_trials = n_trials
        self.pruner = pruner
        self.timeout = timeout

    @abstractmethod
    def suggest(self, trial: Trial) -> Dict[str, Any]:  
        raise NotImplementedError

class DistributionStep(BaseStep):
    def __init__(
        self,
        distributions: Dict[str, distributions.BaseDistribution],
        sampler: samplers.BaseSampler = None,
        pruner: pruners.BasePruner = None,
        n_trials: int = None,
        timeout: int = None,
    ):
    pass   


class GridStep(DistributionStep):
    def __init__(
        self,
        search_space: Dict[str, List[Union[str, float, int, bool, None]]],
        pruner: pruners.BasePruner = None,
        n_trials: int = None,
        timeout: int = None,
    ):
    pass

class StepwiseTuner:
    def __init__(
        self,
        steps: List[Tuple[str, BaseStep]],
        # objective accepts Trial and parameters
        objective: Callable[[optuna.Trial, Dict[str, Any]], float],
        direction: StudyDirection,
        n_trials: int = None,
        timeout: int = None,
        n_jobs: int = 1,
    ):
    pass

    def tune(
        self, storage: Union[None, str, storages.BaseStorage] = None,
    ) -> Tuple[Dict[str, Any], float]:
    # should support callbacks, catch, gc_after_trial arguments?
    pass

Internally, StepwiseTuner.tune() creates a study for each step and an objective function based on Basestep.suggest() (similar to LightGBMTuner implementation). Subclasses of StepwiseTuner, such as XGBoostTuner or LightGBMTuner, deal with the details of the integration. Crucially, they must provide to their parent StepwiseTuner an objective function with the signature Callable[[optuna.Trial, Dict[str, Any]], float]. Basically this objective abstracts away the details of the external library.

StepwiseTuner.tune() returns the best params and best values but those should be class attributes to follow LightGBMTuner interface.

Usage example:

...

params = {
    "seed": 42,
    "objective": "multi:softprob",
    "tree_method": "hist",
    "eta": 0.05,
    "max_depth": 12,
    "max_leaves": 4095,
    "min_child_weight": 1,
    "colsample_bytree": 1,
    "subsample": 0.8,
    "colsample_bytree": 0.8,
    "colsample_bylevel": 0.5,
    "verbosity": 1,
    "num_class": ref_set.num_class,
}

steps = [
    (
        "max_depth and max_leaves",
        GridStep(
            {
                "max_depth": list(range(3, 13)),
                "max_leaves": [15, 31, 63, 127, 255, 511, 1023, 2047, 4095],
            }
        ),
    ),
    (
        "sampling",
        DistributionStep(
            {
                "subsample": DiscreteUniformDistribution(0.6, 1.0, 0.05),
                "colsample_bytree": DiscreteUniformDistribution(0.6, 1.0, 0.05),
                "colsample_bylevel": DiscreteUniformDistribution(0.5, 1.0, 0.05),
            }
        ),
    ),
]

tuner = XgboostTuner(
    params,
    dtrain,
    steps,
    obj=_forecast_avg_obj,
    feval=feval,
    n_jobs=-1,
    n_splits=5,
    num_boost_round=5000,
    early_stopping_rounds=25
)

best_params, best_value = tuner.optimize()

For convenience, default steps can be offered to the end-users. For example (xgboost):

def create_xgb_major_step():
    return [
        (
            "max_depth and max_leaves",
            GridStep(
                {
                    "max_depth": list(range(3, 13)),
                    "max_leaves": [15, 31, 63, 127, 255, 511, 1023, 2047, 4095],
                }
            ),
        )
    ]

... 

# default for XGBoostTuner if no steps given
def create_xgb_comprehensive_steps():
    return [
        create_xgb_major_step(),
        create_xgb_sampling_step(),
        create_xgb_hessian_reg_step(),
        create_xgb_reg_step(),
    ]

[jeffzi]

@toshihikoyanase I’ve submitted the PR.