How to Create a Representative Test Set | by Dimitris Poulopoulos | Towards Data Science

Splitting a data set into train and test sets is usually one of the first processing steps in a machine learning pipeline. After this point, there is just one inviolable rule: set aside the test data split and refer to it again only when your model is ready for its final evaluation. But is there a way to be confident that the split we made is the right one? How can we be sure that the test set is representative of our population?
We must be sure that the test set resembles closely the intricacies of our population. The same goes for our validation sets as well.
In this story, we discuss two methods of splitting: random and stratified sampling. We consider their advantages and debate over which one to use under specific circumstances. This story is actually my published notes on the excellent Hands-On Machine Learning with Scikit-Learn, Keras, and TensorFlow: Concepts, Tools, and Techniques to Build Intelligent Systems by Aurélien Géron.
Learning Rate is my weekly newsletter for those who are curious about the world of AI and MLOps. You’ll hear from me every Friday with updates and thoughts on the latest AI news, research, repos and books. Subscribe here!
For this example, we use the California Housing Prices data set. This data set is excellent for introductory tutorials due to its familiar nature, few missing values, modest size and clear goal.
This dataset is based on data from the 1990 California census. The task is to predict the median value of the houses in a specific district, given a set of features (e.g., housing age, total rooms, etc.). The target or dependent variable (i.e., the median housing value) is a numeric value. Hence, we have a regression task to deal with.
The features of the data set are pretty much self-explanatory, but let us load the CSV file using pandas and look at the first five rows.
We can execute more methods to get a feel for about the data, like df.info() or df.describe(), while a simple histogram of each numerical instance — using the df.hist() method — would give us a better picture of the information in our hands. But, since this is not the scope of this article, let us move onto the next step, which is splitting the data set into train and test sets.
Creating a test set is really easy in theory; just select a few rows at random — typically the 10% or 20% of the full data set — and keep them in a new, separate dataframe. Moreover, for reproducibility, we could set a random_state number so every time we run an experiment we get back the same test set. That would also help compare different approaches directly because every algorithm we choose is evaluated on the same test set.
Moving forward, we use a convenient method provided by Scikit-Learn named train_test_split(). The method accepts lists, numpy arrays, scipy sparse matrices or pandas dataframes. We should also provide the split percentage as it is None by default and a random_state number.
We can see that the method returns a tuple, which we can unpack directly into a train and a test set. Since we pass a dataframe to the method, the resulting train and test sets are also dataframes. Keep in mind that by default the method first shuffles the data set and then splits it. To avoid this behaviour we could set the shuffle attribute to False.
Usually, if the data set is large enough (especially relative to the number of attributes) randomly splitting it would suffice. But if that’s not the case, random splitting could introduce significant sampling bias.
train_set, test_set = train_test_split(df, test_size=.2, random_state=42, shuffle=False)
We now have a randomly generated test set to evaluate our model on. But is it any good? Is it representative of the population? Well, it depends. Usually, if the data set is large enough (especially relative to the number of attributes) randomly splitting it would suffice. But if that’s not the case, random splitting could introduce significant sampling bias. So, let us see how to fix this problem.
Imagine we are running a survey in the US that records the education level and median income of a specific state. If we only question people that live in expensive districts we would surely introduce significant bias in our sample. Thus, we need a way to select people of any economic status roughly approximating the corresponding distribution of the population.
In our case, let us first check if there is a variable that correlates highly with our target variable. This is done easily using the pandas.corr() method.
We sort the results for the dependent variable in descending order and we get that median_income has a fairly strong positive correlation with median_house_value. This is surely something to expect.
It seems that splitting the data set in a stratified fashion, using median_income as the class labels would be a good idea. First, let us create a new variable that creates the class labels for this attribute.
pd.cut will create five classes, according to the bins we specify and keep them in a new variable we name income_cat for income category. We are now ready to split our data set in a stratified fashion.
The only change is that we now set the stratify argument of the train_test_split() method with an array-like object that sets the label classes. There is also another, more involved way to do this, specified in the code sample below for legacy purposes.
In any case, let us see how we did. We will examine the ratio for each income category in the full data set, as well as in the test sample.
Let us compare the two results side by side.
We can see that we are very close. The small differences are really insignificant. We now have a test set that follows the distribution of the income category in our population!
Creating a test set is not always that straight-forward. We must be sure that the test set resembles closely the intricacies of our population. The same goes for our validation sets as well. For instance, check the StratifiedKFold() method provided by Scikit-Learn.
In this story, we saw how we can split a data set into train and test sets both randomly and in a stratified fashion. We implemented the corresponding solutions in Python, using the Scikit-Learn library. Finally, we provided the details and advantages for each method and a simple practical rule on when to use each one.
Learning Rate is my weekly newsletter for those who are curious about the world of AI and MLOps. You’ll hear from me every Friday with updates and thoughts on the latest AI news, research, repos and books. Subscribe here!
My name is Dimitris Poulopoulos and I’m a machine learning researcher at BigDataStack and PhD(c) at the University of Piraeus, Greece. I have worked on designing and implementing AI and software solutions for major clients such as the European Commission, Eurostat, IMF, the European Central Bank, OECD, and IKEA. If you are interested in reading more posts about Machine Learning, Deep Learning and Data Science, follow me on Medium, LinkedIn or @james2pl on twitter.