5 Best Ways to Create Test Datasets Using sklearn in Python

Rate this post

πŸ’‘ Problem Formulation: When building machine learning models, having a well-structured test dataset is critical for evaluating performance. This article explains how to create test datasets in Python using scikit-learn, a powerful machine learning library. Each method below will provide insights into the creation of various types of datasets, suited for different kinds of machine learning problems. For instance, input could be specific parameters for a dataset, and the desired output is a test dataset shaped according to those parameters.

Method 1: Using train_test_split() Function

This method involves splitting a dataset into random train and test subsets using the train_test_split() function. Ideal for when you have an existing dataset and need to evaluate the performance of your machine learning model. The function allows you to specify the proportion of the dataset to include in the test split, and can shuffle the dataset before splitting.

Here’s an example:

from sklearn.model_selection import train_test_split
X, y = [[0, 1], [2, 3], [4, 5], [6, 7]], [0, 1, 2, 3]
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=42)
print(X_test, y_test)


([[4, 5], [0, 1]], [2, 0])

This code snippet demonstrates how to split a simple dataset into training and testing subsets. We first define our features X and target y. Then, we use train_test_split() specifying a test size of 25% of the data, and set a random state for reproducibility. The result is a random sample of our original dataset reserved for testing.

Method 2: Generating Synthetic Classification Data with make_classification()

Sklearn’s make_classification() is a powerful method for generating a random n-class classification problem. This function is helpful when you want to simulate a dataset with controllable noise and a number of informative features. It’s suitable for testing classifiers and visualizing decision boundaries.

Here’s an example:

from sklearn.datasets import make_classification
X, y = make_classification(n_samples=100, n_features=2, n_informative=2, n_redundant=0, random_state=42)
print(X[:5], y[:5])


([[ 0.487, -1.83 ],
 [ 1.24,   1.55 ],
 [-1.36,   1.72 ],
 [ 0.92,   1.16 ],
 [-0.58,  -0.63 ]],
[1, 0, 0, 0, 1])

The example code generates a synthetic binary classification dataset with 100 samples and 2 features. We ensure all features are informative with no redundant features and set a random state for reproducibility. The output shows the first five samples of the feature matrix X along with the corresponding class labels y.

Method 3: Creating Clustered Data Using make_blobs()

For tasks that require testing clustering algorithms, sklearn’s make_blobs() function offers a way to create multi-class datasets by generating isotropic Gaussian blobs. It allows control over the number of features, centers, and cluster standard deviation, which is ideal for evaluating clustering models.

Here’s an example:

from sklearn.datasets import make_blobs
X, y = make_blobs(n_samples=100, centers=3, n_features=2, random_state=42)
print(X[:5], y[:5])


([[ 4.938,  2.985],
 [-6.341,  5.105],
 [-5.122,  4.381],
 [-2.669,  8.815],
 [ 3.575,  1.973]],
[0, 1, 1, 2, 0])

The code above generates 100 samples with 2 features, grouped into 3 clusters. We set a random state to ensure the same data points are generated each time for consistency. The output is a set of feature vectors suitable for clustering, along with the cluster labels.

Method 4: Generating Regression Data with make_regression()

When the need arises for synthetic regression data, the make_regression() function comes into play. This method creates a random regression problem, providing options to specify the number of samples, features, and informative features. It’s instrumental for simulation and testing regression algorithms.

Here’s an example:

from sklearn.datasets import make_regression
X, y = make_regression(n_samples=100, n_features=1, n_informative=1, noise=0.1, random_state=42)
print(X[:5], y[:5])


([[ 0.931],
 [ 0.287],
 [ 0.736],
[-23.74,  -2.05,   1.65,  18.77,  -0.31])

This snippet creates a dataset suited for regression analysis. We have 100 samples, one feature, and have introduced a small amount of noise. The seed for the random number generator is set to obtain reproducible results. The output lists the inputs X with their corresponding continuous target values y.

Bonus One-Liner Method 5: make_data() Function Composition

If you’re looking for a swift inline approach to generate synthetic datasets, composing make_*() functions is your go-to method. This approach is highly adaptable, allowing for quick experimentation and testing various machine learning algorithms on the fly.

Here’s an example:

X, y = make_classification(n_features=4, random_state=0)



Unlike previous examples, this one-liner leverages the defaults of the make_classification() function, immediately creating a dataset with 4 features for a classification problem. It’s concise and perfect for quickly getting a dataset ready for exploratory analysis or the early stages of model development.


  • Method 1: train_test_split(). Strengths: straightforward splitting of existing datasets; customizable test size. Weaknesses: requires initial dataset.
  • Method 2: make_classification(). Strengths: control over noise and the number of informative features; ideal for classification problems. Weaknesses: synthetic data may not represent real-world complexities.
  • Method 3: make_blobs(). Strengths: excellent for clustering problems; control over cluster characteristics. Weaknesses: mainly applicable to clustering problems with well-separated groups.
  • Method 4: make_regression(). Strengths: customizable for regression problem specifics; includes noise parameters. Weaknesses: may oversimplify real-world regression tasks.
  • Method 5: Inline Function Composition. Strengths: quick, one-liner dataset generation. Weaknesses: less customizable and may require further tweaking.