XGBoost案例之红酒品质分类-模型训练

from sklearn.utils.class_weight import compute_sample_weight
sample_weight = compute_sample_weight(class_weight='balanced', y=y_train)

import xgboost as xgb

xgb.XGBClassifier()

import joblib
joblib.dump(model, 'model.pkl')

import pandas as pd
import numpy as np
import xgboost as xgb
from sklearn.utils.class_weight import compute_sample_weight
from sklearn.metrics import accuracy_score, classification_report
import joblib

# 1. 加载数据
def load_data():
    """加载训练集和测试集"""
    print("=== 加载数据 ===")
    train_data = pd.read_csv('data/wine_train.csv')
    test_data = pd.read_csv('data/wine_test.csv')
    
    print(f"训练集: {train_data.shape}")
    print(f"测试集: {test_data.shape}")
    
    return train_data, test_data

# 2. 提取特征和标签
def extract_features_labels(train_data, test_data):
    """提取特征和标签"""
    print("\n=== 提取特征和标签 ===")
    # 最后一列是标签，其他是特征
    X_train = train_data.iloc[:, :-1]
    y_train = train_data.iloc[:, -1]
    
    X_test = test_data.iloc[:, :-1]
    y_test = test_data.iloc[:, -1]
    
    print(f"训练集特征: {X_train.shape}, 标签: {y_train.shape}")
    print(f"测试集特征: {X_test.shape}, 标签: {y_test.shape}")
    
    return X_train, X_test, y_train, y_test

# 3. 平衡权重
def balance_weights(y_train):
    """平衡权重（处理样本不均衡）"""
    print("\n=== 平衡权重 ===")
    sample_weight = compute_sample_weight(
        class_weight='balanced',
        y=y_train
    )
    
    print(f"样本权重形状: {sample_weight.shape}")
    print(f"样本权重范围: {sample_weight.min():.4f} - {sample_weight.max():.4f}")
    
    # 查看各类别权重
    unique_classes = np.unique(y_train)
    for cls in unique_classes:
        mask = y_train == cls
        avg_weight = np.mean(sample_weight[mask])
        count = np.sum(mask)
        print(f"  类别{cls}: 样本数={count}, 平均权重={avg_weight:.4f}")
    
    return sample_weight

# 4. 创建和训练模型
def train_model(X_train, y_train, sample_weight):
    """创建和训练模型"""
    print("\n=== 创建和训练模型 ===")
    
    # 创建XGBoost分类器
    model = xgb.XGBClassifier(
        n_estimators=100,
        max_depth=6,
        learning_rate=0.1,
        subsample=0.8,
        colsample_bytree=0.8,
        objective='multi:softprob',
        num_class=6,
        random_state=22
    )
    
    # 训练模型（传入样本权重）
    model.fit(X_train, y_train, sample_weight=sample_weight)
    
    print("模型训练完成")
    return model

# 5. 评估模型
def evaluate_model(model, X_train, X_test, y_train, y_test):
    """评估模型"""
    print("\n=== 评估模型 ===")
    
    # 训练集评估
    y_train_pred = model.predict(X_train)
    train_acc = accuracy_score(y_train, y_train_pred)
    print(f"训练集准确率: {train_acc:.4f}")
    
    # 测试集评估
    y_test_pred = model.predict(X_test)
    test_acc = accuracy_score(y_test, y_test_pred)
    print(f"测试集准确率: {test_acc:.4f}")
    
    # 分类报告
    print("\n分类报告:")
    print(classification_report(y_test, y_test_pred))
    
    return test_acc

# 6. 保存模型
def save_model(model, path='model/xgboost_wine_model.pkl'):
    """保存模型"""
    print("\n=== 保存模型 ===")
    import os
    os.makedirs('model', exist_ok=True)
    
    joblib.dump(model, path)
    print(f"模型已保存到: {path}")

# 7. 主函数
def main():
    """主函数"""
    print("=" * 50)
    print("XGBoost - 红酒品质分类（模型训练）")
    print("=" * 50)
    
    # 1. 加载数据
    train_data, test_data = load_data()
    
    # 2. 提取特征和标签
    X_train, X_test, y_train, y_test = extract_features_labels(train_data, test_data)
    
    # 3. 平衡权重
    sample_weight = balance_weights(y_train)
    
    # 4. 创建和训练模型
    model = train_model(X_train, y_train, sample_weight)
    
    # 5. 评估模型
    test_acc = evaluate_model(model, X_train, X_test, y_train, y_test)
    
    # 6. 保存模型
    save_model(model)
    
    print("\n" + "=" * 50)
    print("模型训练完成!")
    print("=" * 50)

if __name__ == '__main__':
    main()

from sklearn.utils.class_weight import compute_sample_weight
sample_weight = compute_sample_weight(class_weight='balanced', y=y_train)

model.fit(X_train, y_train, sample_weight=sample_weight)