I have a dataset for a binary classification task which has 90 percent 'yes' and 10 percent 'no'. Let's say I want to take 25 percent of the data as the test set (which the model will not see). How can I do that in such a way that the class proportion is preserved. For example, if I take 10000 rows as the test data, I want 9000 rows to have class 'yes' and 1000 rows to have class 'no'.
$\begingroup$
$\endgroup$
Add a comment
|
1 Answer
$\begingroup$
$\endgroup$
If you're going to use library functions, you'll use the stratify option in scikit-learn train_test_split. It's not well written how to use it in the documentation, but you'll input your class labels for stratify parameter, e.g.
train_test_split(X, y, stratify = y, test_ratio = 0.25)
If you want to write it from scratch, you can sample from each class directly and combine them to form the test set, i.e. sample 0.25 of class 1 and class 0, and combine them to obtain a 0.25 sample of the entire training set. However, train_test_split does it for your automatically.
