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

Question

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

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

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

Answer 1

你有 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)

Answer 2

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