Dealing with the test set of imbalanced data
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()


&gt; A    60
&gt; B    30
&gt; C    20
&gt; 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()


&gt; A    58
&gt; B    28
&gt; C    18
&gt; Name: category, dtype: int64


test_df.category.value_counts()


&gt; A    2
&gt; B    2
&gt; C    2
&gt; Name: category, dtype: int64