Toggle Menu
Prev Up Next

scikit-learn 1.3.2
Other versions

Please cite us if you use the software.

# `sklearn.model_selection`.StratifiedKFold

class sklearn.model_selection.StratifiedKFold(n_splits=5, *, shuffle=False, random_state=None)

  Stratified K-Folds cross-validator.

  Provides train/test indices to split data in train/test sets.

  This cross-validation object is a variation of KFold that returns stratified folds. The folds are made by preserving the percentage of samples for each class.

  Read more in the User Guide.

  For visualisation of cross-validation behaviour and comparison between common scikit-learn split methods refer to Visualizing cross-validation behavior in scikit-learn

  Parameters:
    n_splits int, default=5

      Number of folds. Must be at least 2.

      Changed in version 0.22: `n_splits` default value changed from 3 to 5.

    shuffle bool, default=False

      Whether to shuffle each class’s samples before splitting into batches. Note that the samples within each split will not be shuffled.

    random_state int, RandomState instance or None, default=None

      When `shuffle` is True, `random_state` affects the ordering of the indices, which controls the randomness of each fold for each class. Otherwise, leave `random_state` as `None`. Pass an int for reproducible output across multiple function calls. See Glossary.

  See also

  `RepeatedStratifiedKFold`

    Repeats Stratified K-Fold n times.

  Notes

  The implementation is designed to:

  - Generate test sets such that all contain the same distribution of classes, or as close as possible.

  - Be invariant to class label: relabelling `y = ["Happy", "Sad"]` to `y = [1, 0]` should not change the indices generated.

  - Preserve order dependencies in the dataset ordering, when `shuffle=False`: all samples from class k in some test set were contiguous in y, or separated in y by samples from classes other than k.

  - Generate test sets where the smallest and largest differ by at most one sample.

  Changed in version 0.22: The previous implementation did not follow the last constraint.

  Examples

      >>> import numpy as np
      >>> from sklearn.model_selection import StratifiedKFold
      >>> X = np.array([[1, 2], [3, 4], [1, 2], [3, 4]])
      >>> y = np.array([0, 0, 1, 1])
      >>> skf = StratifiedKFold(n_splits=2)
      >>> skf.get_n_splits(X, y)
      2
      >>> print(skf)
      StratifiedKFold(n_splits=2, random_state=None, shuffle=False)
      >>> for i, (train_index, test_index) in enumerate(skf.split(X, y)):
      ...     print(f"Fold {i}:")
      ...     print(f"  Train: index={train_index}")
      ...     print(f"  Test:  index={test_index}")
      Fold 0:
        Train: index=[1 3]
        Test:  index=[0 2]
      Fold 1:
        Train: index=[0 2]
        Test:  index=[1 3]

  Methods

  `get_metadata_routing`() | Get metadata routing of this object.
  `get_n_splits`([X, y, groups]) | Returns the number of splitting iterations in the cross-validator
  `split`(X, y[, groups]) | Generate indices to split data into training and test set.

  get_metadata_routing()

    Get metadata routing of this object.

    Please check User Guide on how the routing mechanism works.

    Returns:
      routing MetadataRequest

        A `MetadataRequest` encapsulating routing information.

  get_n_splits(X=None, y=None, groups=None)

    Returns the number of splitting iterations in the cross-validator

    Parameters:
      X object

        Always ignored, exists for compatibility.

      y object

        Always ignored, exists for compatibility.

      groups object

        Always ignored, exists for compatibility.

    Returns:
      n_splits int

        Returns the number of splitting iterations in the cross-validator.

  split(X, y, groups=None)

    Generate indices to split data into training and test set.

    Parameters:
      X array-like of shape (n_samples, n_features)

        Training data, where `n_samples` is the number of samples and `n_features` is the number of features.

        Note that providing `y` is sufficient to generate the splits and hence `np.zeros(n_samples)` may be used as a placeholder for `X` instead of actual training data.

      y array-like of shape (n_samples,)

        The target variable for supervised learning problems. Stratification is done based on the y labels.

      groups object

        Always ignored, exists for compatibility.

    Yields:
      train ndarray

        The training set indices for that split.

      test ndarray

        The testing set indices for that split.

    Notes

    Randomized CV splitters may return different results for each call of split. You can make the results identical by setting `random_state` to an integer.

## Examples using `sklearn.model_selection.StratifiedKFold`

Recursive feature elimination with cross-validation

Recursive feature elimination with cross-validation

GMM covariances

GMM covariances

Receiver Operating Characteristic (ROC) with cross validation

Receiver Operating Characteristic (ROC) with cross validation

Test with permutations the significance of a classification score

Test with permutations the significance of a classification score

Visualizing cross-validation behavior in scikit-learn

Visualizing cross-validation behavior in scikit-learn

Effect of varying threshold for self-training

Effect of varying threshold for self-training