0
$\begingroup$

I am working on a problem dealing with unbalanced data that has a very specific request.

I would like to know the following:

When I have an imbalanced dataset and I do train test split, the test samples also have an imbalance. How can I at least balance the test set so that when I do the confusion matrix equal numbers are tested?

Here's an example (Classification):

If I have three groups: A 60,B 30,C 20. For simplicity, the test set is 10%. So, instead of having a train/test set of 54/6, 27/3, 18/2, I would like 58/2, 28/2 and 18/2 OR 54/2, 27/2, 18/2.

How can I do this?

$\endgroup$
2
$\begingroup$

For a case like this, it would probably be the easiest to do the train/test split manually. Here is a quick example with Pandas:

import pandas as pd
import numpy as np

# Create an empty dataframe with 60*A, 30*B, 20*C and some random target values
df = pd.DataFrame(columns=['category'])
df['category'] = (["A"] * 60 + ["B"] * 30 + ["C"] * 20)
df['target'] = np.random.rand(len(df))
df.category.value_counts()
> A    60
> B    30
> C    20
> Name: category, dtype: int64
# Obtain the size of the least represented category
smallest_category_size = np.min(df.category.value_counts().values)

# Size of the test split of the smallest category
test_size = 0.1

# Test count (2 in this case)
test_count = int(test_size*smallest_category_size)

# New train/test dataframes
train_df = pd.DataFrame(columns=df.columns)
test_df = pd.DataFrame(columns=df.columns)

for cat in df.category.unique():
    # Select only a part of the dataset containing one category
    cat_df = df[df.category == cat]

    # Create and shuffle boolean mask
    mask = np.array((len(cat_df) - test_count) * [True] + test_count * [False])
    np.random.shuffle(mask)

    # Append the selected values to the train/test
    train_df = train_df.append(cat_df[mask])
    test_df = test_df.append(cat_df[~mask])
train_df.category.value_counts()
> A    58
> B    28
> C    18
> Name: category, dtype: int64
test_df.category.value_counts()
> A    2
> B    2
> C    2
> Name: category, dtype: int64
| improve this answer | |
$\endgroup$
1
$\begingroup$

I would still use scikit-learn's built-in mechanism for train/test splitting but downsample after.

import numpy as np
from sklearn.model_selection import train_test_split

# Create mock data
n = 110
X = np.arange(n).reshape((n, 1))
y = [0]*60 + [1]*30 + [2]*20 # Uneven categories

# Standard train test split
X_train, X_test, y_train, y_test = train_test_split(X, y, stratify=y) 

# Fixed number of samples from each category
n_test = 2

# A category sampling
a_indices = [i for i, n in enumerate(y_test) if n == 0]
a_indices_sub_sample = np.random.choice(a_indices, n_test, replace=False)
X_test_a_category = X_test[a_indices_sub_sample]

# B category sampling
b_indices = [i for i, n in enumerate(y_test) if n == 1]
b_indices_sub_sample = np.random.choice(b_indices, n_test, replace=False)
X_test_b_category = X_test[b_indices_sub_sample]

# C category sampling
c_indices = [i for i, n in enumerate(y_test) if n == 2]
c_indices_sub_sample = np.random.choice(c_indices, n_test, replace=False)
X_test_c_category = X_test[c_indices_sub_sample]
| improve this answer | |
$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

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