engine.py 33.9 KB
Newer Older
wxchan's avatar
wxchan committed
1
# coding: utf-8
2
"""Library with training routines of LightGBM."""
3
import copy
4
import json
5
from collections import OrderedDict, defaultdict
wxchan's avatar
wxchan committed
6
from operator import attrgetter
7
from pathlib import Path
8
from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
9

wxchan's avatar
wxchan committed
10
import numpy as np
11

wxchan's avatar
wxchan committed
12
from . import callback
13
from .basic import (Booster, Dataset, LightGBMError, _choose_param_value, _ConfigAliases, _InnerPredictor,
14
15
16
                    _LGBM_BoosterEvalMethodResultType, _LGBM_CategoricalFeatureConfiguration,
                    _LGBM_CustomObjectiveFunction, _LGBM_EvalFunctionResultType, _LGBM_FeatureNameConfiguration,
                    _log_warning)
17
from .compat import SKLEARN_INSTALLED, _LGBMBaseCrossValidator, _LGBMGroupKFold, _LGBMStratifiedKFold
wxchan's avatar
wxchan committed
18

19
20
21
22
23
24
25
__all__ = [
    'cv',
    'CVBooster',
    'train',
]


26
27
28
29
30
31
32
33
34
_LGBM_CustomMetricFunction = Union[
    Callable[
        [np.ndarray, Dataset],
        _LGBM_EvalFunctionResultType,
    ],
    Callable[
        [np.ndarray, Dataset],
        List[_LGBM_EvalFunctionResultType]
    ],
35
]
wxchan's avatar
wxchan committed
36

37
38
39
40
41
_LGBM_PreprocFunction = Callable[
    [Dataset, Dataset, Dict[str, Any]],
    Tuple[Dataset, Dataset, Dict[str, Any]]
]

42
43
44
45
46
47
48
49
50

def train(
    params: Dict[str, Any],
    train_set: Dataset,
    num_boost_round: int = 100,
    valid_sets: Optional[List[Dataset]] = None,
    valid_names: Optional[List[str]] = None,
    feval: Optional[Union[_LGBM_CustomMetricFunction, List[_LGBM_CustomMetricFunction]]] = None,
    init_model: Optional[Union[str, Path, Booster]] = None,
51
52
    feature_name: _LGBM_FeatureNameConfiguration = 'auto',
    categorical_feature: _LGBM_CategoricalFeatureConfiguration = 'auto',
53
54
55
    keep_training_booster: bool = False,
    callbacks: Optional[List[Callable]] = None
) -> Booster:
56
    """Perform the training with given parameters.
wxchan's avatar
wxchan committed
57
58
59
60

    Parameters
    ----------
    params : dict
61
62
        Parameters for training. Values passed through ``params`` take precedence over those
        supplied via arguments.
Guolin Ke's avatar
Guolin Ke committed
63
    train_set : Dataset
64
65
        Data to be trained on.
    num_boost_round : int, optional (default=100)
wxchan's avatar
wxchan committed
66
        Number of boosting iterations.
67
    valid_sets : list of Dataset, or None, optional (default=None)
68
        List of data to be evaluated on during training.
69
    valid_names : list of str, or None, optional (default=None)
70
        Names of ``valid_sets``.
71
    feval : callable, list of callable, or None, optional (default=None)
wxchan's avatar
wxchan committed
72
        Customized evaluation function.
Akshita Dixit's avatar
Akshita Dixit committed
73
        Each evaluation function should accept two parameters: preds, eval_data,
74
        and return (eval_name, eval_result, is_higher_better) or list of such tuples.
75

76
            preds : numpy 1-D array or numpy 2-D array (for multi-class task)
77
                The predicted values.
78
                For multi-class task, preds are numpy 2-D array of shape = [n_samples, n_classes].
79
                If custom objective function is used, predicted values are returned before any transformation,
80
                e.g. they are raw margin instead of probability of positive class for binary task in this case.
Akshita Dixit's avatar
Akshita Dixit committed
81
            eval_data : Dataset
82
                A ``Dataset`` to evaluate.
83
            eval_name : str
84
                The name of evaluation function (without whitespaces).
85
86
87
88
89
            eval_result : float
                The eval result.
            is_higher_better : bool
                Is eval result higher better, e.g. AUC is ``is_higher_better``.

90
91
        To ignore the default metric corresponding to the used objective,
        set the ``metric`` parameter to the string ``"None"`` in ``params``.
92
    init_model : str, pathlib.Path, Booster or None, optional (default=None)
93
        Filename of LightGBM model or Booster instance used for continue training.
94
    feature_name : list of str, or 'auto', optional (default="auto")
95
96
        Feature names.
        If 'auto' and data is pandas DataFrame, data columns names are used.
97
    categorical_feature : list of str or int, or 'auto', optional (default="auto")
98
99
        Categorical features.
        If list of int, interpreted as indices.
100
        If list of str, interpreted as feature names (need to specify ``feature_name`` as well).
101
        If 'auto' and data is pandas DataFrame, pandas unordered categorical columns are used.
102
        All values in categorical features will be cast to int32 and thus should be less than int32 max value (2147483647).
103
        Large values could be memory consuming. Consider using consecutive integers starting from zero.
104
        All negative values in categorical features will be treated as missing values.
105
        The output cannot be monotonically constrained with respect to a categorical feature.
106
        Floating point numbers in categorical features will be rounded towards 0.
107
108
109
    keep_training_booster : bool, optional (default=False)
        Whether the returned Booster will be used to keep training.
        If False, the returned value will be converted into _InnerPredictor before returning.
110
        This means you won't be able to use ``eval``, ``eval_train`` or ``eval_valid`` methods of the returned Booster.
111
112
        When your model is very large and cause the memory error,
        you can try to set this param to ``True`` to avoid the model conversion performed during the internal call of ``model_to_string``.
113
        You can still use _InnerPredictor as ``init_model`` for future continue training.
114
    callbacks : list of callable, or None, optional (default=None)
115
        List of callback functions that are applied at each iteration.
116
        See Callbacks in Python API for more information.
wxchan's avatar
wxchan committed
117

118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
    Note
    ----
    A custom objective function can be provided for the ``objective`` parameter.
    It should accept two parameters: preds, train_data and return (grad, hess).

        preds : numpy 1-D array or numpy 2-D array (for multi-class task)
            The predicted values.
            Predicted values are returned before any transformation,
            e.g. they are raw margin instead of probability of positive class for binary task.
        train_data : Dataset
            The training dataset.
        grad : numpy 1-D array or numpy 2-D array (for multi-class task)
            The value of the first order derivative (gradient) of the loss
            with respect to the elements of preds for each sample point.
        hess : numpy 1-D array or numpy 2-D array (for multi-class task)
            The value of the second order derivative (Hessian) of the loss
            with respect to the elements of preds for each sample point.

    For multi-class task, preds are numpy 2-D array of shape = [n_samples, n_classes],
    and grad and hess should be returned in the same format.

wxchan's avatar
wxchan committed
139
140
    Returns
    -------
141
142
    booster : Booster
        The trained Booster model.
wxchan's avatar
wxchan committed
143
    """
144
145
146
147
148
149
150
151
152
153
154
155
156
157
    if not isinstance(train_set, Dataset):
        raise TypeError(f"train() only accepts Dataset object, train_set has type '{type(train_set).__name__}'.")

    if num_boost_round <= 0:
        raise ValueError(f"num_boost_round must be greater than 0. Got {num_boost_round}.")

    if isinstance(valid_sets, list):
        for i, valid_item in enumerate(valid_sets):
            if not isinstance(valid_item, Dataset):
                raise TypeError(
                    "Every item in valid_sets must be a Dataset object. "
                    f"Item {i} has type '{type(valid_item).__name__}'."
                )

158
    # create predictor first
159
    params = copy.deepcopy(params)
160
161
162
163
164
    params = _choose_param_value(
        main_param_name='objective',
        params=params,
        default_value=None
    )
165
    fobj: Optional[_LGBM_CustomObjectiveFunction] = None
166
167
168
    if callable(params["objective"]):
        fobj = params["objective"]
        params["objective"] = 'none'
169
    for alias in _ConfigAliases.get("num_iterations"):
170
        if alias in params:
171
            num_boost_round = params.pop(alias)
172
            _log_warning(f"Found `{alias}` in params. Will use it instead of argument")
173
    params["num_iterations"] = num_boost_round
174
175
176
177
178
179
180
181
    # setting early stopping via global params should be possible
    params = _choose_param_value(
        main_param_name="early_stopping_round",
        params=params,
        default_value=None
    )
    if params["early_stopping_round"] is None:
        params.pop("early_stopping_round")
182
    first_metric_only = params.get('first_metric_only', False)
183

184
    predictor: Optional[_InnerPredictor] = None
185
    if isinstance(init_model, (str, Path)):
186
187
188
189
        predictor = _InnerPredictor.from_model_file(
            model_file=init_model,
            pred_parameter=params
        )
wxchan's avatar
wxchan committed
190
    elif isinstance(init_model, Booster):
191
192
193
194
195
196
197
198
199
        predictor = _InnerPredictor.from_booster(
            booster=init_model,
            pred_parameter=dict(init_model.params, **params)
        )

    if predictor is not None:
        init_iteration = predictor.current_iteration()
    else:
        init_iteration = 0
Guolin Ke's avatar
Guolin Ke committed
200

201
202
203
204
    train_set._update_params(params) \
             ._set_predictor(predictor) \
             .set_feature_name(feature_name) \
             .set_categorical_feature(categorical_feature)
Guolin Ke's avatar
Guolin Ke committed
205

wxchan's avatar
wxchan committed
206
207
    is_valid_contain_train = False
    train_data_name = "training"
Guolin Ke's avatar
Guolin Ke committed
208
    reduced_valid_sets = []
wxchan's avatar
wxchan committed
209
    name_valid_sets = []
210
    if valid_sets is not None:
Guolin Ke's avatar
Guolin Ke committed
211
212
        if isinstance(valid_sets, Dataset):
            valid_sets = [valid_sets]
213
        if isinstance(valid_names, str):
wxchan's avatar
wxchan committed
214
            valid_names = [valid_names]
Guolin Ke's avatar
Guolin Ke committed
215
        for i, valid_data in enumerate(valid_sets):
216
            # reduce cost for prediction training data
Guolin Ke's avatar
Guolin Ke committed
217
            if valid_data is train_set:
wxchan's avatar
wxchan committed
218
219
220
221
                is_valid_contain_train = True
                if valid_names is not None:
                    train_data_name = valid_names[i]
                continue
Nikita Titov's avatar
Nikita Titov committed
222
            reduced_valid_sets.append(valid_data._update_params(params).set_reference(train_set))
223
            if valid_names is not None and len(valid_names) > i:
wxchan's avatar
wxchan committed
224
225
                name_valid_sets.append(valid_names[i])
            else:
226
                name_valid_sets.append(f'valid_{i}')
227
    # process callbacks
228
    if callbacks is None:
229
        callbacks_set = set()
wxchan's avatar
wxchan committed
230
231
232
    else:
        for i, cb in enumerate(callbacks):
            cb.__dict__.setdefault('order', i - len(callbacks))
233
        callbacks_set = set(callbacks)
wxchan's avatar
wxchan committed
234

235
236
237
    if "early_stopping_round" in params:
        callbacks_set.add(
            callback.early_stopping(
238
                stopping_rounds=params["early_stopping_round"],  # type: ignore[arg-type]
239
240
241
242
243
244
245
246
                first_metric_only=first_metric_only,
                verbose=_choose_param_value(
                    main_param_name="verbosity",
                    params=params,
                    default_value=1
                ).pop("verbosity") > 0
            )
        )
247

248
249
250
251
    callbacks_before_iter_set = {cb for cb in callbacks_set if getattr(cb, 'before_iteration', False)}
    callbacks_after_iter_set = callbacks_set - callbacks_before_iter_set
    callbacks_before_iter = sorted(callbacks_before_iter_set, key=attrgetter('order'))
    callbacks_after_iter = sorted(callbacks_after_iter_set, key=attrgetter('order'))
wxchan's avatar
wxchan committed
252

253
    # construct booster
254
255
256
257
    try:
        booster = Booster(params=params, train_set=train_set)
        if is_valid_contain_train:
            booster.set_train_data_name(train_data_name)
258
        for valid_set, name_valid_set in zip(reduced_valid_sets, name_valid_sets):
259
260
261
262
263
            booster.add_valid(valid_set, name_valid_set)
    finally:
        train_set._reverse_update_params()
        for valid_set in reduced_valid_sets:
            valid_set._reverse_update_params()
264
    booster.best_iteration = 0
wxchan's avatar
wxchan committed
265

266
    # start training
267
    for i in range(init_iteration, init_iteration + num_boost_round):
wxchan's avatar
wxchan committed
268
269
        for cb in callbacks_before_iter:
            cb(callback.CallbackEnv(model=booster,
270
                                    params=params,
wxchan's avatar
wxchan committed
271
                                    iteration=i,
272
273
                                    begin_iteration=init_iteration,
                                    end_iteration=init_iteration + num_boost_round,
wxchan's avatar
wxchan committed
274
275
276
277
                                    evaluation_result_list=None))

        booster.update(fobj=fobj)

278
        evaluation_result_list: List[_LGBM_BoosterEvalMethodResultType] = []
wxchan's avatar
wxchan committed
279
        # check evaluation result.
280
        if valid_sets is not None:
wxchan's avatar
wxchan committed
281
282
283
284
285
286
            if is_valid_contain_train:
                evaluation_result_list.extend(booster.eval_train(feval))
            evaluation_result_list.extend(booster.eval_valid(feval))
        try:
            for cb in callbacks_after_iter:
                cb(callback.CallbackEnv(model=booster,
287
                                        params=params,
wxchan's avatar
wxchan committed
288
                                        iteration=i,
289
290
                                        begin_iteration=init_iteration,
                                        end_iteration=init_iteration + num_boost_round,
wxchan's avatar
wxchan committed
291
                                        evaluation_result_list=evaluation_result_list))
292
293
        except callback.EarlyStopException as earlyStopException:
            booster.best_iteration = earlyStopException.best_iteration + 1
wxchan's avatar
wxchan committed
294
            evaluation_result_list = earlyStopException.best_score
wxchan's avatar
wxchan committed
295
            break
296
    booster.best_score = defaultdict(OrderedDict)
wxchan's avatar
wxchan committed
297
298
    for dataset_name, eval_name, score, _ in evaluation_result_list:
        booster.best_score[dataset_name][eval_name] = score
299
    if not keep_training_booster:
300
        booster.model_from_string(booster.model_to_string()).free_dataset()
wxchan's avatar
wxchan committed
301
302
303
    return booster


304
class CVBooster:
305
306
    """CVBooster in LightGBM.

307
    Auxiliary data structure to hold and redirect all boosters of ``cv()`` function.
308
    This class has the same methods as Booster class.
309
310
311
312
313
314
    All method calls, except for the following methods, are actually performed for underlying Boosters and
    then all returned results are returned in a list.

    - ``model_from_string()``
    - ``model_to_string()``
    - ``save_model()``
315
316
317
318
319
320
321
322

    Attributes
    ----------
    boosters : list of Booster
        The list of underlying fitted models.
    best_iteration : int
        The best iteration of fitted model.
    """
323

324
325
326
327
    def __init__(
        self,
        model_file: Optional[Union[str, Path]] = None
    ):
328
329
        """Initialize the CVBooster.

330
331
332
333
        Parameters
        ----------
        model_file : str, pathlib.Path or None, optional (default=None)
            Path to the CVBooster model file.
334
        """
335
        self.boosters: List[Booster] = []
336
        self.best_iteration = -1
337

338
339
340
341
        if model_file is not None:
            with open(model_file, "r") as file:
                self._from_dict(json.load(file))

342
    def _append(self, booster: Booster) -> None:
343
        """Add a booster to CVBooster."""
344
345
        self.boosters.append(booster)

346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
    def _from_dict(self, models: Dict[str, Any]) -> None:
        """Load CVBooster from dict."""
        self.best_iteration = models["best_iteration"]
        self.boosters = []
        for model_str in models["boosters"]:
            self._append(Booster(model_str=model_str))

    def _to_dict(self, num_iteration: Optional[int], start_iteration: int, importance_type: str) -> Dict[str, Any]:
        """Serialize CVBooster to dict."""
        models_str = []
        for booster in self.boosters:
            models_str.append(booster.model_to_string(num_iteration=num_iteration, start_iteration=start_iteration,
                                                      importance_type=importance_type))
        return {"boosters": models_str, "best_iteration": self.best_iteration}

361
    def __getattr__(self, name: str) -> Callable[[Any, Any], List[Any]]:
362
        """Redirect methods call of CVBooster."""
363
        def handler_function(*args: Any, **kwargs: Any) -> List[Any]:
364
            """Call methods with each booster, and concatenate their results."""
365
366
367
368
            ret = []
            for booster in self.boosters:
                ret.append(getattr(booster, name)(*args, **kwargs))
            return ret
369
        return handler_function
wxchan's avatar
wxchan committed
370

371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
    def __getstate__(self) -> Dict[str, Any]:
        return vars(self)

    def __setstate__(self, state: Dict[str, Any]) -> None:
        vars(self).update(state)

    def model_from_string(self, model_str: str) -> "CVBooster":
        """Load CVBooster from a string.

        Parameters
        ----------
        model_str : str
            Model will be loaded from this string.

        Returns
        -------
        self : CVBooster
            Loaded CVBooster object.
        """
        self._from_dict(json.loads(model_str))
        return self

    def model_to_string(
        self,
        num_iteration: Optional[int] = None,
        start_iteration: int = 0,
        importance_type: str = 'split'
    ) -> str:
        """Save CVBooster to JSON string.

        Parameters
        ----------
        num_iteration : int or None, optional (default=None)
            Index of the iteration that should be saved.
            If None, if the best iteration exists, it is saved; otherwise, all iterations are saved.
            If <= 0, all iterations are saved.
        start_iteration : int, optional (default=0)
            Start index of the iteration that should be saved.
        importance_type : str, optional (default="split")
            What type of feature importance should be saved.
            If "split", result contains numbers of times the feature is used in a model.
            If "gain", result contains total gains of splits which use the feature.

        Returns
        -------
        str_repr : str
            JSON string representation of CVBooster.
        """
        return json.dumps(self._to_dict(num_iteration, start_iteration, importance_type))

    def save_model(
        self,
        filename: Union[str, Path],
        num_iteration: Optional[int] = None,
        start_iteration: int = 0,
        importance_type: str = 'split'
    ) -> "CVBooster":
        """Save CVBooster to a file as JSON text.

        Parameters
        ----------
        filename : str or pathlib.Path
            Filename to save CVBooster.
        num_iteration : int or None, optional (default=None)
            Index of the iteration that should be saved.
            If None, if the best iteration exists, it is saved; otherwise, all iterations are saved.
            If <= 0, all iterations are saved.
        start_iteration : int, optional (default=0)
            Start index of the iteration that should be saved.
        importance_type : str, optional (default="split")
            What type of feature importance should be saved.
            If "split", result contains numbers of times the feature is used in a model.
            If "gain", result contains total gains of splits which use the feature.

        Returns
        -------
        self : CVBooster
            Returns self.
        """
        with open(filename, "w") as file:
            json.dump(self._to_dict(num_iteration, start_iteration, importance_type), file)

        return self

455

456
457
458
459
460
461
def _make_n_folds(
    full_data: Dataset,
    folds: Optional[Union[Iterable[Tuple[np.ndarray, np.ndarray]], _LGBMBaseCrossValidator]],
    nfold: int,
    params: Dict[str, Any],
    seed: int,
462
463
464
465
    fpreproc: Optional[_LGBM_PreprocFunction],
    stratified: bool,
    shuffle: bool,
    eval_train_metric: bool
466
) -> CVBooster:
467
    """Make a n-fold list of Booster from random indices."""
wxchan's avatar
wxchan committed
468
469
    full_data = full_data.construct()
    num_data = full_data.num_data()
470
    if folds is not None:
471
472
473
474
475
476
        if not hasattr(folds, '__iter__') and not hasattr(folds, 'split'):
            raise AttributeError("folds should be a generator or iterator of (train_idx, test_idx) tuples "
                                 "or scikit-learn splitter object with split method")
        if hasattr(folds, 'split'):
            group_info = full_data.get_group()
            if group_info is not None:
477
                group_info = np.array(group_info, dtype=np.int32, copy=False)
478
                flatted_group = np.repeat(range(len(group_info)), repeats=group_info)
479
            else:
480
                flatted_group = np.zeros(num_data, dtype=np.int32)
481
            folds = folds.split(X=np.empty(num_data), y=full_data.get_label(), groups=flatted_group)
wxchan's avatar
wxchan committed
482
    else:
483
484
485
        if any(params.get(obj_alias, "") in {"lambdarank", "rank_xendcg", "xendcg",
                                             "xe_ndcg", "xe_ndcg_mart", "xendcg_mart"}
               for obj_alias in _ConfigAliases.get("objective")):
wxchan's avatar
wxchan committed
486
            if not SKLEARN_INSTALLED:
487
                raise LightGBMError('scikit-learn is required for ranking cv')
488
            # ranking task, split according to groups
489
            group_info = np.array(full_data.get_group(), dtype=np.int32, copy=False)
490
            flatted_group = np.repeat(range(len(group_info)), repeats=group_info)
491
            group_kfold = _LGBMGroupKFold(n_splits=nfold)
492
            folds = group_kfold.split(X=np.empty(num_data), groups=flatted_group)
wxchan's avatar
wxchan committed
493
494
        elif stratified:
            if not SKLEARN_INSTALLED:
495
                raise LightGBMError('scikit-learn is required for stratified cv')
496
            skf = _LGBMStratifiedKFold(n_splits=nfold, shuffle=shuffle, random_state=seed)
497
            folds = skf.split(X=np.empty(num_data), y=full_data.get_label())
extremin's avatar
extremin committed
498
        else:
wxchan's avatar
wxchan committed
499
500
501
502
503
            if shuffle:
                randidx = np.random.RandomState(seed).permutation(num_data)
            else:
                randidx = np.arange(num_data)
            kstep = int(num_data / nfold)
504
505
506
            test_id = [randidx[i: i + kstep] for i in range(0, num_data, kstep)]
            train_id = [np.concatenate([test_id[i] for i in range(nfold) if k != i]) for k in range(nfold)]
            folds = zip(train_id, test_id)
wxchan's avatar
wxchan committed
507

508
    ret = CVBooster()
wxchan's avatar
wxchan committed
509
    for train_idx, test_idx in folds:
510
511
        train_set = full_data.subset(sorted(train_idx))
        valid_set = full_data.subset(sorted(test_idx))
wxchan's avatar
wxchan committed
512
513
        # run preprocessing on the data set if needed
        if fpreproc is not None:
wxchan's avatar
wxchan committed
514
            train_set, valid_set, tparam = fpreproc(train_set, valid_set, params.copy())
wxchan's avatar
wxchan committed
515
        else:
wxchan's avatar
wxchan committed
516
            tparam = params
517
        cvbooster = Booster(tparam, train_set)
518
519
        if eval_train_metric:
            cvbooster.add_valid(train_set, 'train')
520
        cvbooster.add_valid(valid_set, 'valid')
521
        ret._append(cvbooster)
wxchan's avatar
wxchan committed
522
523
    return ret

wxchan's avatar
wxchan committed
524

525
526
527
def _agg_cv_result(
    raw_results: List[List[Tuple[str, str, float, bool]]]
) -> List[Tuple[str, str, float, bool, float]]:
528
    """Aggregate cross-validation results."""
529
    cvmap: Dict[str, List[float]] = OrderedDict()
530
    metric_type: Dict[str, bool] = {}
wxchan's avatar
wxchan committed
531
532
    for one_result in raw_results:
        for one_line in one_result:
533
            key = f"{one_line[0]} {one_line[1]}"
534
            metric_type[key] = one_line[3]
535
            cvmap.setdefault(key, [])
536
            cvmap[key].append(one_line[2])
wxchan's avatar
wxchan committed
537
    return [('cv_agg', k, np.mean(v), metric_type[k], np.std(v)) for k, v in cvmap.items()]
wxchan's avatar
wxchan committed
538

wxchan's avatar
wxchan committed
539

540
541
542
543
544
545
546
547
548
549
550
def cv(
    params: Dict[str, Any],
    train_set: Dataset,
    num_boost_round: int = 100,
    folds: Optional[Union[Iterable[Tuple[np.ndarray, np.ndarray]], _LGBMBaseCrossValidator]] = None,
    nfold: int = 5,
    stratified: bool = True,
    shuffle: bool = True,
    metrics: Optional[Union[str, List[str]]] = None,
    feval: Optional[Union[_LGBM_CustomMetricFunction, List[_LGBM_CustomMetricFunction]]] = None,
    init_model: Optional[Union[str, Path, Booster]] = None,
551
552
    feature_name: _LGBM_FeatureNameConfiguration = 'auto',
    categorical_feature: _LGBM_CategoricalFeatureConfiguration = 'auto',
553
554
555
556
557
    fpreproc: Optional[_LGBM_PreprocFunction] = None,
    seed: int = 0,
    callbacks: Optional[List[Callable]] = None,
    eval_train_metric: bool = False,
    return_cvbooster: bool = False
558
) -> Dict[str, Union[List[float], CVBooster]]:
Andrew Ziem's avatar
Andrew Ziem committed
559
    """Perform the cross-validation with given parameters.
wxchan's avatar
wxchan committed
560
561
562
563

    Parameters
    ----------
    params : dict
564
565
        Parameters for training. Values passed through ``params`` take precedence over those
        supplied via arguments.
Guolin Ke's avatar
Guolin Ke committed
566
    train_set : Dataset
567
        Data to be trained on.
568
    num_boost_round : int, optional (default=100)
wxchan's avatar
wxchan committed
569
        Number of boosting iterations.
570
    folds : generator or iterator of (train_idx, test_idx) tuples, scikit-learn splitter object or None, optional (default=None)
571
        If generator or iterator, it should yield the train and test indices for each fold.
572
        If object, it should be one of the scikit-learn splitter classes
573
        (https://scikit-learn.org/stable/modules/classes.html#splitter-classes)
574
        and have ``split`` method.
575
        This argument has highest priority over other data split arguments.
576
    nfold : int, optional (default=5)
wxchan's avatar
wxchan committed
577
        Number of folds in CV.
578
579
    stratified : bool, optional (default=True)
        Whether to perform stratified sampling.
580
    shuffle : bool, optional (default=True)
581
        Whether to shuffle before splitting data.
582
    metrics : str, list of str, or None, optional (default=None)
583
584
        Evaluation metrics to be monitored while CV.
        If not None, the metric in ``params`` will be overridden.
585
    feval : callable, list of callable, or None, optional (default=None)
586
        Customized evaluation function.
587
        Each evaluation function should accept two parameters: preds, eval_data,
588
        and return (eval_name, eval_result, is_higher_better) or list of such tuples.
589

590
            preds : numpy 1-D array or numpy 2-D array (for multi-class task)
591
                The predicted values.
592
                For multi-class task, preds are numpy 2-D array of shape = [n_samples, n_classes].
593
                If custom objective function is used, predicted values are returned before any transformation,
594
                e.g. they are raw margin instead of probability of positive class for binary task in this case.
595
596
            eval_data : Dataset
                A ``Dataset`` to evaluate.
597
            eval_name : str
Andrew Ziem's avatar
Andrew Ziem committed
598
                The name of evaluation function (without whitespace).
599
600
601
602
603
            eval_result : float
                The eval result.
            is_higher_better : bool
                Is eval result higher better, e.g. AUC is ``is_higher_better``.

604
605
        To ignore the default metric corresponding to the used objective,
        set ``metrics`` to the string ``"None"``.
606
    init_model : str, pathlib.Path, Booster or None, optional (default=None)
607
        Filename of LightGBM model or Booster instance used for continue training.
608
    feature_name : list of str, or 'auto', optional (default="auto")
609
610
        Feature names.
        If 'auto' and data is pandas DataFrame, data columns names are used.
611
    categorical_feature : list of str or int, or 'auto', optional (default="auto")
612
613
        Categorical features.
        If list of int, interpreted as indices.
614
        If list of str, interpreted as feature names (need to specify ``feature_name`` as well).
615
        If 'auto' and data is pandas DataFrame, pandas unordered categorical columns are used.
616
        All values in categorical features will be cast to int32 and thus should be less than int32 max value (2147483647).
617
        Large values could be memory consuming. Consider using consecutive integers starting from zero.
618
        All negative values in categorical features will be treated as missing values.
619
        The output cannot be monotonically constrained with respect to a categorical feature.
620
        Floating point numbers in categorical features will be rounded towards 0.
621
622
    fpreproc : callable or None, optional (default=None)
        Preprocessing function that takes (dtrain, dtest, params)
wxchan's avatar
wxchan committed
623
        and returns transformed versions of those.
624
    seed : int, optional (default=0)
wxchan's avatar
wxchan committed
625
        Seed used to generate the folds (passed to numpy.random.seed).
626
    callbacks : list of callable, or None, optional (default=None)
627
        List of callback functions that are applied at each iteration.
628
        See Callbacks in Python API for more information.
629
630
631
    eval_train_metric : bool, optional (default=False)
        Whether to display the train metric in progress.
        The score of the metric is calculated again after each training step, so there is some impact on performance.
632
633
    return_cvbooster : bool, optional (default=False)
        Whether to return Booster models trained on each fold through ``CVBooster``.
wxchan's avatar
wxchan committed
634

635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
    Note
    ----
    A custom objective function can be provided for the ``objective`` parameter.
    It should accept two parameters: preds, train_data and return (grad, hess).

        preds : numpy 1-D array or numpy 2-D array (for multi-class task)
            The predicted values.
            Predicted values are returned before any transformation,
            e.g. they are raw margin instead of probability of positive class for binary task.
        train_data : Dataset
            The training dataset.
        grad : numpy 1-D array or numpy 2-D array (for multi-class task)
            The value of the first order derivative (gradient) of the loss
            with respect to the elements of preds for each sample point.
        hess : numpy 1-D array or numpy 2-D array (for multi-class task)
            The value of the second order derivative (Hessian) of the loss
            with respect to the elements of preds for each sample point.

    For multi-class task, preds are numpy 2-D array of shape = [n_samples, n_classes],
    and grad and hess should be returned in the same format.

wxchan's avatar
wxchan committed
656
657
    Returns
    -------
658
659
660
661
    eval_hist : dict
        Evaluation history.
        The dictionary has the following format:
        {'metric1-mean': [values], 'metric1-stdv': [values],
Qiwei Ye's avatar
Qiwei Ye committed
662
        'metric2-mean': [values], 'metric2-stdv': [values],
663
        ...}.
664
        If ``return_cvbooster=True``, also returns trained boosters wrapped in a ``CVBooster`` object via ``cvbooster`` key.
wxchan's avatar
wxchan committed
665
    """
Guolin Ke's avatar
Guolin Ke committed
666
    if not isinstance(train_set, Dataset):
667
668
669
670
671
        raise TypeError(f"cv() only accepts Dataset object, train_set has type '{type(train_set).__name__}'.")

    if num_boost_round <= 0:
        raise ValueError(f"num_boost_round must be greater than 0. Got {num_boost_round}.")

672
    params = copy.deepcopy(params)
673
674
675
676
677
    params = _choose_param_value(
        main_param_name='objective',
        params=params,
        default_value=None
    )
678
    fobj: Optional[_LGBM_CustomObjectiveFunction] = None
679
680
681
    if callable(params["objective"]):
        fobj = params["objective"]
        params["objective"] = 'none'
682
    for alias in _ConfigAliases.get("num_iterations"):
683
        if alias in params:
684
            _log_warning(f"Found '{alias}' in params. Will use it instead of 'num_boost_round' argument")
685
            num_boost_round = params.pop(alias)
686
    params["num_iterations"] = num_boost_round
687
688
689
690
691
692
693
694
    # setting early stopping via global params should be possible
    params = _choose_param_value(
        main_param_name="early_stopping_round",
        params=params,
        default_value=None
    )
    if params["early_stopping_round"] is None:
        params.pop("early_stopping_round")
695
    first_metric_only = params.get('first_metric_only', False)
696

697
    if isinstance(init_model, (str, Path)):
698
699
700
701
        predictor = _InnerPredictor.from_model_file(
            model_file=init_model,
            pred_parameter=params
        )
Guolin Ke's avatar
Guolin Ke committed
702
    elif isinstance(init_model, Booster):
703
704
705
706
        predictor = _InnerPredictor.from_booster(
            booster=init_model,
            pred_parameter=dict(init_model.params, **params)
        )
Guolin Ke's avatar
Guolin Ke committed
707
708
709
    else:
        predictor = None

Peter's avatar
Peter committed
710
    if metrics is not None:
711
712
        for metric_alias in _ConfigAliases.get("metric"):
            params.pop(metric_alias, None)
Peter's avatar
Peter committed
713
        params['metric'] = metrics
wxchan's avatar
wxchan committed
714

715
716
717
718
719
    train_set._update_params(params) \
             ._set_predictor(predictor) \
             .set_feature_name(feature_name) \
             .set_categorical_feature(categorical_feature)

720
    results = defaultdict(list)
721
    cvfolds = _make_n_folds(full_data=train_set, folds=folds, nfold=nfold,
722
                            params=params, seed=seed, fpreproc=fpreproc,
723
724
                            stratified=stratified, shuffle=shuffle,
                            eval_train_metric=eval_train_metric)
wxchan's avatar
wxchan committed
725
726

    # setup callbacks
727
    if callbacks is None:
728
        callbacks_set = set()
wxchan's avatar
wxchan committed
729
730
731
    else:
        for i, cb in enumerate(callbacks):
            cb.__dict__.setdefault('order', i - len(callbacks))
732
        callbacks_set = set(callbacks)
733
734

    if "early_stopping_round" in params:
735
        callbacks_set.add(
736
            callback.early_stopping(
737
                stopping_rounds=params["early_stopping_round"],  # type: ignore[arg-type]
738
739
740
741
742
743
744
745
                first_metric_only=first_metric_only,
                verbose=_choose_param_value(
                    main_param_name="verbosity",
                    params=params,
                    default_value=1
                ).pop("verbosity") > 0
            )
        )
wxchan's avatar
wxchan committed
746

747
748
749
750
    callbacks_before_iter_set = {cb for cb in callbacks_set if getattr(cb, 'before_iteration', False)}
    callbacks_after_iter_set = callbacks_set - callbacks_before_iter_set
    callbacks_before_iter = sorted(callbacks_before_iter_set, key=attrgetter('order'))
    callbacks_after_iter = sorted(callbacks_after_iter_set, key=attrgetter('order'))
wxchan's avatar
wxchan committed
751

752
    for i in range(num_boost_round):
wxchan's avatar
wxchan committed
753
        for cb in callbacks_before_iter:
754
755
            cb(callback.CallbackEnv(model=cvfolds,
                                    params=params,
wxchan's avatar
wxchan committed
756
757
758
759
                                    iteration=i,
                                    begin_iteration=0,
                                    end_iteration=num_boost_round,
                                    evaluation_result_list=None))
760
761
        cvfolds.update(fobj=fobj)  # type: ignore[call-arg]
        res = _agg_cv_result(cvfolds.eval_valid(feval))  # type: ignore[call-arg]
wxchan's avatar
wxchan committed
762
        for _, key, mean, _, std in res:
763
764
            results[f'{key}-mean'].append(mean)
            results[f'{key}-stdv'].append(std)
wxchan's avatar
wxchan committed
765
766
        try:
            for cb in callbacks_after_iter:
767
768
                cb(callback.CallbackEnv(model=cvfolds,
                                        params=params,
wxchan's avatar
wxchan committed
769
770
771
772
                                        iteration=i,
                                        begin_iteration=0,
                                        end_iteration=num_boost_round,
                                        evaluation_result_list=res))
773
774
        except callback.EarlyStopException as earlyStopException:
            cvfolds.best_iteration = earlyStopException.best_iteration + 1
775
776
            for bst in cvfolds.boosters:
                bst.best_iteration = cvfolds.best_iteration
wxchan's avatar
wxchan committed
777
            for k in results:
778
                results[k] = results[k][:cvfolds.best_iteration]
wxchan's avatar
wxchan committed
779
            break
780
781

    if return_cvbooster:
782
        results['cvbooster'] = cvfolds  # type: ignore[assignment]
783

wxchan's avatar
wxchan committed
784
    return dict(results)