我使用以下代码将我的数据集拆分为训练/验证/测试集。
from sklearn.model_selection import train_test_splitX_train, X_test, y_train, y_test = train_test_split(X_data, y_data, test_size=0.3, random_state=42)X_test, X_val, y_test, y_val = train_test_split(X_test, y_test, test_size=0.5, random_state=42)
问题是我的数据集非常不平衡。例如,某些类别有500个样本,而有些只有70个。在这种情况下,这种拆分方法准确吗?抽样是随机的,还是 sklearn 使用了一些方法来保持数据在所有集合中的分布相同?
回答:
你应该使用 stratify
选项(参见文档):
X_train, X_test, y_train, y_test = train_test_split(X_data, y_data, test_size=0.3, random_state=42, stratify=y_data)X_test, X_val, y_test, y_val = train_test_split(X_test, y_test, test_size=0.5, random_state=42, stratify=y_test)