GBDT算法之泰坦尼克号案例

GradientBoostingClassifier()

GridSearchCV(
    estimator,
    param_grid,
    cv
)

import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.metrics import accuracy_score
from sklearn.feature_extraction import DictVectorizer

# 1. 加载数据
def load_data():
    """加载泰坦尼克号数据"""
    data = pd.read_csv('data/train.csv')
    print(f"数据形状: {data.shape}")
    return data

# 2. 数据预处理
def preprocess_data(data):
    """数据预处理"""
    # 2.1 提取特征和标签
    X = data[['Pclass', 'Age', 'Sex']].copy()
    y = data['Survived']
    
    # 2.2 缺失值填充（Age用均值）
    X['Age'] = X['Age'].fillna(X['Age'].mean())
    
    # 2.3 One-Hot编码（Sex列）
    X_dict = X.to_dict(orient='records')
    transfer = DictVectorizer(sparse=False)
    X_new = transfer.fit_transform(X_dict)
    
    print(f"特征形状: {X_new.shape}")
    print(f"特征名称: {transfer.get_feature_names_out()}")
    
    return X_new, y, transfer

# 3. 划分数据集
def split_data(X, y):
    """划分训练集和测试集"""
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=22
    )
    print(f"训练集: {X_train.shape}")
    print(f"测试集: {X_test.shape}")
    return X_train, X_test, y_train, y_test

# 4. GBDT基本训练
def train_gbdt(X_train, X_test, y_train, y_test):
    """GBDT基本训练"""
    print("\n=== GBDT基本训练 ===")
    gbc = GradientBoostingClassifier(random_state=22)
    gbc.fit(X_train, y_train)
    
    y_pred = gbc.predict(X_test)
    acc = accuracy_score(y_test, y_pred)
    print(f"准确率: {acc:.4f}")
    return gbc, acc

# 5. 网格搜索调参
def grid_search(X_train, X_test, y_train, y_test):
    """网格搜索调参"""
    print("\n=== 网格搜索调参 ===")
    gbc = GradientBoostingClassifier(random_state=22)
    
    param_grid = {
        'n_estimators': [50, 100, 200],
        'learning_rate': [0.01, 0.1, 0.2],
        'max_depth': [3, 5, 7]
    }
    
    grid_search = GridSearchCV(
        gbc,
        param_grid,
        cv=3,
        n_jobs=-1
    )
    
    grid_search.fit(X_train, y_train)
    
    print(f"最优参数: {grid_search.best_params_}")
    print(f"最优得分: {grid_search.best_score_:.4f}")
    
    best_model = grid_search.best_estimator_
    y_pred = best_model.predict(X_test)
    acc = accuracy_score(y_test, y_pred)
    print(f"测试集准确率: {acc:.4f}")
    
    return best_model, acc

# 6. 主函数
def main():
    """主函数"""
    print("=" * 50)
    print("GBDT - 泰坦尼克号案例")
    print("=" * 50)
    
    # 1. 加载数据
    data = load_data()
    
    # 2. 数据预处理
    X, y, transfer = preprocess_data(data)
    
    # 3. 划分数据集
    X_train, X_test, y_train, y_test = split_data(X, y)
    
    # 4. GBDT基本训练
    gbc, gbc_acc = train_gbdt(X_train, X_test, y_train, y_test)
    
    # 5. 网格搜索调参
    best_model, best_acc = grid_search(X_train, X_test, y_train, y_test)
    
    # 6. 对比
    print("\n=== 对比 ===")
    print(f"GBDT基本: {gbc_acc:.4f}")
    print(f"GBDT调参后: {best_acc:.4f}")
    
    # 7. 预测新样本
    print("\n=== 预测新样本 ===")
    samples = [
        {'Pclass': 3, 'Age': 25.0, 'Sex': 'male'},
        {'Pclass': 1, 'Age': 30.0, 'Sex': 'female'},
        {'Pclass': 2, 'Age': 35.0, 'Sex': 'male'}
    ]
    
    X_new = transfer.transform(samples)
    predictions = best_model.predict(X_new)
    
    for sample, pred in zip(samples, predictions):
        status = '存活' if pred == 1 else '未存活'
        print(f"{sample}: {status}")

if __name__ == '__main__':
    main()