在Python中，划分数据集为训练集和测试集是机器学习中的常见步骤，主要目的是为了评估模型的泛化能力。以下是几种常见的实现方法。 ### 1. 使用 `train_test_split` 方法 `scikit-learn` 库中的 `train_test_split` 函数是划分数据集最简单且常用的方法之一。它可以直接将数据集按照指定比例随机划分为训练集和测试集。 ```python from sklearn.model_selection import train_test_split # 假设 X 是特征数据，y 是标签数据 X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) ``` - `test_size=0.2` 表示测试集占总数据的20%。 - `random_state=42` 确保每次划分的结果一致。 ### 2. 使用 K 折交叉验证 K 折交叉验证（K-Fold Cross Validation）是一种更稳健的评估方法，它将数据集划分为 K 个子集（称为“折”），每次使用其中一折作为测试集，其余 K-1 折作为训练集，重复 K 次并取平均结果。 ```python from sklearn.model_selection import KFold kf = KFold(n_splits=5, shuffle=True, random_state=42) for train_index, test_index in kf.split(X): X_train, X_test = X[train_index], X[test_index] y_train, y_test = y[train_index], y[test_index] ``` - `n_splits=5` 表示将数据集划分为5折。 - `shuffle=True` 表示在划分前先打乱数据。 ### 3. 使用分层抽样（Stratified K-Fold） 当数据集存在类别不平衡问题时，可以使用 `StratifiedKFold` 来保持每一折中各类别的比例与原始数据集一致。 ```python from sklearn.model_selection import StratifiedKFold skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42) for train_index, test_index in skf.split(X, y): X_train, X_test = X[train_index], X[test_index] y_train, y_test = y[train_index], y[test_index] ``` ### 4. 自助法（Bootstrap） 自助法是一种有放回抽样方法，适用于小数据集。它通过从原始数据集中重复抽样构建训练集，未被抽中的样本作为测试集。 ```python import numpy as np def bootstrap_sample(X, y, random_state=42): np.random.seed(random_state) n_samples = len(X) indices = np.random.choice(n_samples, size=n_samples, replace=True) out_of_bag = np.setdiff1d(np.arange(n_samples), indices) return X[indices], y[indices], X[out_of_bag], y[out_of_bag] X_train, y_train, X_test, y_test = bootstrap_sample(X, y) ``` ### 5. 手动划分数据集 如果希望对数据划分过程有更精细的控制，也可以手动实现划分逻辑。 ```python import numpy as np # 随机打乱数据 indices = np.random.permutation(len(X)) train_size = int(0.8 * len(X)) train_indices, test_indices = indices[:train_size], indices[train_size:] X_train, X_test = X[train_indices], X[test_indices] y_train, y_test = y[train_indices], y[test_indices] ``` ###