Python中如何分割数据集为训练集和测试集并建立训练模型？

# 完整示例：使用scikit-learn分割数据集并训练模型
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.datasets import load_iris
from sklearn.metrics import accuracy_score

# 1. 加载示例数据集
data = load_iris()
X = data.data  # 特征
y = data.target  # 标签

# 2. 分割数据集
# test_size: 测试集比例，random_state: 随机种子确保可重复性
X_train, X_test, y_train, y_test = train_test_split(
    X, y, 
    test_size=0.2,  # 20%作为测试集
    random_state=42,  # 固定随机种子
    stratify=y  # 保持类别比例
)

print(f"训练集样本数: {len(X_train)}")
print(f"测试集样本数: {len(X_test)}")

# 3. 建立并训练模型
model = LogisticRegression(max_iter=200)
model.fit(X_train, y_train)

# 4. 预测和评估
y_pred = model.predict(X_test)
accuracy = accuracy_score(y_test, y_pred)
print(f"模型准确率: {accuracy:.2%}")

# 5. 使用模型进行新预测
new_sample = np.array([[5.1, 3.5, 1.4, 0.2]])  # 新样本
prediction = model.predict(new_sample)
print(f"新样本预测类别: {prediction[0]}")
print(f"对应类别名: {data.target_names[prediction[0]]}")

关键点说明：

1. train_test_split是主要的分割函数，参数test_size通常设为0.2-0.3
2. random_state保证每次分割结果一致，便于调试
3. stratify=y在分类问题中很重要，能保持训练集和测试集的类别分布一致
4. 训练模型只需调用fit()方法，预测用predict()

一句话建议： 记得设置随机种子保证结果可复现。