博客

用pandas划分数据集实现训练集和测试集

  • Post category:Python

下面是详细的步骤和示例说明：

1. 准备数据

首先要使用 pandas 导入数据，并进行数据处理以便于划分数据集，例如去除空值、重复值、无用的特征等。

import pandas as pd

# 导入数据
data = pd.read_csv('data.csv')

# 数据预处理
data.dropna()  # 删除空值
data.drop_duplicates()  # 删除重复值
data.drop('id', axis=1, inplace=True)  # 删除无用的特征

2. 划分数据集

接着，使用 train_test_split() 函数将数据集划分为训练集和测试集。该函数需要传递的参数包括数据集 X 和标签 y，以及测试集占总数据的比例 test_size。

from sklearn.model_selection import train_test_split

# 划分数据集
X = data.drop('label', axis=1)  # 特征
y = data['label']  # 标签
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

此时，数据集已经成功划分为训练集和测试集，且比例为 7:3。

3. 计算划分后的数据集大小

划分好数据集后，可以使用 len() 函数计算训练集和测试集的数量，以确保划分比例正确。

# 计算数据集大小
train_size = len(X_train)
test_size = len(X_test)
print(f'Training dataset size: {train_size}')
print(f'Testing dataset size: {test_size}')

输出结果如下：

Training dataset size: 7000
Testing dataset size: 3000

4. 示例说明1

以鸢尾花数据集为例，使用 pandas 划分数据集实现训练集和测试集：

import pandas as pd
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split

# 加载数据集
iris = load_iris()

# 转化为 pandas DataFrame
data = pd.DataFrame(iris.data, columns=iris.feature_names)
data['label'] = iris.target

# 划分数据集
X = data.drop('label', axis=1)
y = data['label']
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

# 打印数据集大小
print(f'Training dataset size: {len(X_train)}')
print(f'Testing dataset size: {len(X_test)}')

输出结果如下：

Training dataset size: 105
Testing dataset size: 45

5. 示例说明2

以鸢尾花数据集为例，使用 pandas 划分数据集实现训练集和测试集，并使用决策树模型进行分类：

import pandas as pd
from sklearn.datasets import load_iris
from sklearn.tree import DecisionTreeClassifier
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split

# 加载数据集
iris = load_iris()

# 转化为 pandas DataFrame
data = pd.DataFrame(iris.data, columns=iris.feature_names)
data['label'] = iris.target

# 划分数据集
X = data.drop('label', axis=1)
y = data['label']
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

# 训练模型
model = DecisionTreeClassifier()
model.fit(X_train, y_train)

# 预测测试集
y_pred = model.predict(X_test)

# 计算分类准确率
accuracy = accuracy_score(y_test, y_pred)
print(f'Accuracy: {accuracy}')

输出结果如下：

Accuracy: 0.9777777777777777

以上就是使用 pandas 划分数据集实现训练集和测试集的完整攻略，同时给出了两个示例说明，希望对您有所帮助。

Tags: