I'm reading through the book Hands-On Machine Learning with Scikit-Learn and Tensorflow by Aurélien Géron. In a regression project on California Housing Prices, he goes over the concept of stratified sampling.

I think I understand the concept as his explanation "the population is divided into homogeneous subgroups called strata, and the right number of instances is sampled from each stratum to guarantee that the test set is representative of the overall population."

So in my own words, simply splitting the dataset with sklearn's train_test_split leaves the train and test set vulnerable to misrepresenting the ratios of categorical variables (ie population has 40% one category, 60% another, but the train/test set are totally different ratios of these categories), so stratifying ensures the sample is 'random', but still maintaining proper ratios within test and train splits. Please correct me if I'm wrong.

Here's the code to his stratified sampling based on income categories (housing is the main dataframe):

split = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=42)
for train_index, test_index in split.split(housing, housing['income_cat']):
    strat_train_set = housing.loc[train_index]
    strat_test_set = housing.loc[test_index]

I'm pretty confused with this code:

1) What does the variable 'split' represent? Does it comprise both the train and test split...?

2) In the 2nd line of code, what does split.split mean? I guess I'm confused with most of how StratifiedShuffleSplit divides the train and test set and why he needed to create this 'for' loop in order to create strat_train_set and strat_test_set.



  • $\begingroup$ I also got confused at the first reading. The code would have been much better(less confusing) if the object is named differently, say stratified_sample = StratifiedShuffleSplit(...). $\endgroup$ Jan 17, 2020 at 15:36

3 Answers 3


Is the object that allows us to do stratified split, and split.split is the split’s object method/function called split that can be used to perform stratified split.

Depending on how many splits you want using the n_split parameter you can split, hence the for loop.

  • $\begingroup$ Gotcha, so first split is var I created called split, and second split is method/function to perform stratified split on var split. So what is var split without using its split method/function? What am I creating in the first line of code and what am I iterating through with my for loop? $\endgroup$
    – Greg Rosen
    Apr 30, 2019 at 6:03
  • 2
    $\begingroup$ You are creating an object that needs to be initialized first to be used. This is how sklearn does it’s API designs. You create an object with specified params and then use it later as supposed to a function that you can directly use. $\endgroup$ Apr 30, 2019 at 6:08

The 'StratifiedShuffleSplit' function takes parameters on how the split needs to take place and returns a function to do the split.

The 'split' variable in the first line is used to store this function. In Python, functions/procedures can be stored as variables.

'n_splits' indicates the number of folds. 'test_size' indicates the proportion of the test data set out of the complete dataset.

The for loop allows you to iterate over the multiple splits when n_splits>1. Example, if n_splits=2 and test_size=0.2, the dataset 'housing' is split into two sets of (80% training and 20% test)

@Greg Rosen


I too had similar doubts as OP, then I went through the Scikit-Learn's GitHub repo. Repo Link to the file in discussion

And found out that:

  1. split.split(housing, housing['income_cat']): Is nothing but the first split is the object of class StratifiedShuffleSplit() i.e can be any name you want to keep, and the second split is the function in this class (that refers to split function of base class BaseShuffleSplit()) the use of this is to Generate indices to split data into training and test set.

  2. The reason for using a loop is best explained by @SparkZeus above.

Happy building models !!


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.