在每个分组中训练/测试分组保留类的比例

问题描述

假设我想拆分如下所示的NER数据集：

Data: "Jokowi is the president of Indonesia"
Label: ['B-Person','O','Country']

是否有任何python库或算法可以确保训练数据集和测试数据集的每个类分布相同？ 任何建议将不胜感激

解决方法

您可以使用sklearn的StratifiedShuffleSplit来做到这一点。从文档中：

折叠是通过保留每个类别的样本百分比来完成的。

StratifiedShuffleSplit返回一个生成器，其中包含将数据帧分为训练和测试的索引。这是一个示例用例，可以清楚地表明，每个拆分中确实保留了类比例：

from sklearn.model_selection import StratifiedShuffleSplit
import seaborn as sns

X = np.random.randint(0,5,(1200,2))
y = np.random.choice([0,1],size=(1200,),p=[0.8,0.2])

sss = StratifiedShuffleSplit(n_splits=2,test_size=0.2,random_state=0)
train_index,test_index = next(sss.split(X,y))

fig,axes = plt.subplots(1,2,figsize=(10,5))
for split,title,ax in zip([train_index,test_index],['Train split','Test split'],axes.flatten()):
    sns.countplot(y[split],ax=ax).set_title(title)

,

您可以探索Scikit学习库中提供的StratifiedShuffleSplit。