Skip to content

train_xgboost

def train_xgboost(config, budget, X, y, seed=None) #

Train an XGBoost model on the given data.

PARAMETER DESCRIPTION
config

The configuration to use for the XGBoost model

TYPE: Configuration

budget

The number of estimators to use for the XGBoost model

TYPE: int

X

The data to train on

TYPE: DataFrame

y

The target to train on

TYPE: Series

seed

The seed to use for the XGBoost model

TYPE: int | None DEFAULT: None

RETURNS DESCRIPTION
XGBRegressor

The trained XGBoost model
Source code in src/mfpbench/pd1/surrogate/train_xgboost.py 
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
def train_xgboost(
    config: Configuration,
    budget: int,
    X: pd.DataFrame,
    y: pd.Series,
    seed: int | None = None,
) -> XGBRegressor:
    """Train an XGBoost model on the given data.

    Args:
        config: The configuration to use for the XGBoost model
        budget: The number of estimators to use for the XGBoost model
        X: The data to train on
        y: The target to train on
        seed: The seed to use for the XGBoost model

    Returns:
        The trained XGBoost model
    """
    from xgboost import XGBRegressor

    if y.name == "train_cost":
        model = XGBRegressor(
            **config,
            seed=seed,
            n_estimators=budget,
            monotone_constraints={"epoch": 1},
        )
    else:
        model = XGBRegressor(**config, seed=seed, n_estimators=budget)

    model.fit(X, y)

    return model  # type: ignore