step2_classification.py

# -*- coding: utf-8 -*-
"""
水质安全预测项目 - 步骤2: 构建监督分类任务
目的: 学习如何将实际问题转化为机器学习问题，构建和训练分类模型
"""

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset
import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
import joblib
import warnings

warnings.filterwarnings('ignore')

# 设置中文字体显示
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


def load_preprocessed_data(data_dir="preprocessed_data"):
    """
    加载预处理后的数据
    从步骤1的输出中加载清理过的数据
    """
    print("=" * 50)
    print("步骤2: 加载预处理数据")
    print("=" * 50)

    try:
        # 加载PyTorch张量
        X_tensor = torch.load(f"{data_dir}/X_tensor.pt")
        y_tensor = torch.load(f"{data_dir}/y_tensor.pt")

        # 加载标准化器
        scaler = joblib.load(f"{data_dir}/scaler.pkl")

        print(f"✓ 成功加载预处理数据")
        print(f"特征数据形状: {X_tensor.shape}")
        print(f"标签数据形状: {y_tensor.shape}")
        print(f"特征数量: {X_tensor.shape[1]}")
        print(f"样本数量: {X_tensor.shape[0]}")

        # 检查类别分布
        class_counts = torch.bincount(y_tensor)
        print(f"\n类别分布:")
        print(f"  不可饮用 (0): {class_counts[0]} 样本 ({class_counts[0] / len(y_tensor) * 100:.1f}%)")
        print(f"  可饮用 (1): {class_counts[1]} 样本 ({class_counts[1] / len(y_tensor) * 100:.1f}%)")

        return X_tensor, y_tensor, scaler

    except Exception as e:
        print(f"❌ 数据加载失败: {e}")
        print("请先运行步骤1进行数据预处理")
        return None, None, None


def split_dataset(X, y, test_size=0.2, val_size=0.2, random_state=42):
    """
    划分数据集为训练集、验证集和测试集
    学习机器学习中正确的数据划分策略
    """
    print("\n" + "=" * 50)
    print("步骤2.1: 数据集划分")
    print("=" * 50)

    print(f"数据划分策略:")
    print(f"  训练集: {(1 - test_size - val_size) * 100:.0f}% (用于模型训练)")
    print(f"  验证集: {val_size * 100:.0f}% (用于模型调优和早停)")
    print(f"  测试集: {test_size * 100:.0f}% (用于最终性能评估)")

    # 转换为numpy进行划分
    X_np = X.numpy()
    y_np = y.numpy()

    # 首先分离出测试集
    X_temp, X_test, y_temp, y_test = train_test_split(
        X_np, y_np, test_size=test_size, random_state=random_state, stratify=y_np
    )

    # 再从剩余数据中分离出验证集
    val_size_adjusted = val_size / (1 - test_size)  # 调整验证集比例
    X_train, X_val, y_train, y_val = train_test_split(
        X_temp, y_temp, test_size=val_size_adjusted, random_state=random_state, stratify=y_temp
    )

    # 转换回PyTorch张量
    X_train = torch.FloatTensor(X_train)
    X_val = torch.FloatTensor(X_val)
    X_test = torch.FloatTensor(X_test)
    y_train = torch.LongTensor(y_train)
    y_val = torch.LongTensor(y_val)
    y_test = torch.LongTensor(y_test)

    print(f"\n数据集划分结果:")
    print(f"  训练集: {X_train.shape[0]} 样本")
    print(f"  验证集: {X_val.shape[0]} 样本")
    print(f"  测试集: {X_test.shape[0]} 样本")

    # 验证类别平衡性
    print(f"\n各数据集的类别分布:")
    for name, y_split in [("训练集", y_train), ("验证集", y_val), ("测试集", y_test)]:
        class_counts = torch.bincount(y_split)
        pos_ratio = class_counts[1] / len(y_split)
        print(f"  {name}: 可饮用比例 = {pos_ratio:.3f}")

    return X_train, X_val, X_test, y_train, y_val, y_test


def create_data_loaders(X_train, y_train, X_val, y_val, X_test, y_test, batch_size=64):
    """
    创建PyTorch数据加载器
    学习如何高效地批量加载数据进行训练
    """
    print("\n" + "=" * 50)
    print("步骤2.2: 创建数据加载器")
    print("=" * 50)

    print(f"批处理大小 (Batch Size): {batch_size}")
    print(f"为什么使用批处理:")
    print(f"  1. 内存效率 - 不需要一次加载所有数据")
    print(f"  2. 训练稳定性 - 每批数据的梯度更新更稳定")
    print(f"  3. 计算效率 - 利用GPU并行计算能力")

    # 创建数据集
    train_dataset = TensorDataset(X_train, y_train)
    val_dataset = TensorDataset(X_val, y_val)
    test_dataset = TensorDataset(X_test, y_test)

    # 创建数据加载器
    train_loader = DataLoader(
        train_dataset,
        batch_size=batch_size,
        shuffle=True,  # 训练时打乱数据顺序
        drop_last=False  # 保留最后一个不完整的批次
    )

    val_loader = DataLoader(
        val_dataset,
        batch_size=batch_size,
        shuffle=False,  # 验证时不需要打乱
        drop_last=False
    )

    test_loader = DataLoader(
        test_dataset,
        batch_size=batch_size,
        shuffle=False,  # 测试时不需要打乱
        drop_last=False
    )

    print(f"\n数据加载器信息:")
    print(f"  训练批次数: {len(train_loader)}")
    print(f"  验证批次数: {len(val_loader)}")
    print(f"  测试批次数: {len(test_loader)}")

    return train_loader, val_loader, test_loader


class WaterQualityClassifier(nn.Module):
    """
    水质分类神经网络模型
    学习如何设计适合分类任务的神经网络架构
    """

    def __init__(self, input_size, hidden_size=64, dropout_rate=0.3):
        """
        初始化模型架构

        Args:
            input_size: 输入特征数量 (9个水质参数)
            hidden_size: 隐藏层神经元数量
            dropout_rate: Dropout率，用于防止过拟合
        """
        super(WaterQualityClassifier, self).__init__()

        print(f"构建神经网络模型:")
        print(f"  输入层: {input_size} 个特征")
        print(f"  隐藏层1: {hidden_size} 个神经元")
        print(f"  隐藏层2: {hidden_size // 2} 个神经元")
        print(f"  输出层: 2 个类别 (可饮用/不可饮用)")
        print(f"  Dropout率: {dropout_rate} (防止过拟合)")

        # 定义网络层
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.dropout1 = nn.Dropout(dropout_rate)

        self.fc2 = nn.Linear(hidden_size, hidden_size // 2)
        self.dropout2 = nn.Dropout(dropout_rate)

        self.fc3 = nn.Linear(hidden_size // 2, 2)  # 2个输出类别

        # 激活函数
        self.relu = nn.ReLU()

    def forward(self, x):
        """
        前向传播过程
        定义数据如何在网络中流动
        """
        # 第一个隐藏层
        x = self.fc1(x)
        x = self.relu(x)
        x = self.dropout1(x)

        # 第二个隐藏层
        x = self.fc2(x)
        x = self.relu(x)
        x = self.dropout2(x)

        # 输出层
        x = self.fc3(x)

        return x


def train_model(model, train_loader, val_loader, num_epochs=100, learning_rate=0.001, patience=10):
    """
    训练分类模型
    学习监督学习的核心训练过程
    """
    print("\n" + "=" * 50)
    print("步骤2.3: 模型训练")
    print("=" * 50)

    # 设置训练设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    print(f"训练设备: {device}")

    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()  # 分类任务使用交叉熵损失
    optimizer = optim.Adam(model.parameters(), lr=learning_rate, weight_decay=1e-4)

    print(f"训练超参数:")
    print(f"  学习率: {learning_rate}")
    print(f"  最大训练轮数: {num_epochs}")
    print(f"  早停耐心值: {patience} (验证损失不改善时提前停止)")
    print(f"  损失函数: 交叉熵损失 (适合分类任务)")
    print(f"  优化器: Adam (自适应学习率)")

    # 记录训练过程
    train_losses = []
    val_losses = []
    train_accuracies = []
    val_accuracies = []

    # 早停相关变量
    best_val_loss = float('inf')
    patience_counter = 0
    best_model_state = None

    print(f"\n开始训练...")
    print(f"{'轮次':<6} {'训练损失':<10} {'验证损失':<10} {'训练准确率':<12} {'验证准确率':<12}")
    print("-" * 60)

    for epoch in range(num_epochs):
        # 训练阶段
        model.train()
        train_loss = 0.0
        correct_train = 0
        total_train = 0

        for batch_X, batch_y in train_loader:
            batch_X, batch_y = batch_X.to(device), batch_y.to(device)

            # 前向传播
            outputs = model(batch_X)
            loss = criterion(outputs, batch_y)

            # 反向传播和优化
            optimizer.zero_grad()  # 清零梯度
            loss.backward()  # 计算梯度
            optimizer.step()  # 更新参数

            # 记录训练统计
            train_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            total_train += batch_y.size(0)
            correct_train += (predicted == batch_y).sum().item()

        # 计算平均训练损失和准确率
        avg_train_loss = train_loss / len(train_loader)
        train_accuracy = correct_train / total_train

        # 验证阶段
        model.eval()
        val_loss = 0.0
        correct_val = 0
        total_val = 0

        with torch.no_grad():  # 验证时不计算梯度
            for batch_X, batch_y in val_loader:
                batch_X, batch_y = batch_X.to(device), batch_y.to(device)

                outputs = model(batch_X)
                loss = criterion(outputs, batch_y)

                val_loss += loss.item()
                _, predicted = torch.max(outputs.data, 1)
                total_val += batch_y.size(0)
                correct_val += (predicted == batch_y).sum().item()

        # 计算平均验证损失和准确率
        avg_val_loss = val_loss / len(val_loader)
        val_accuracy = correct_val / total_val

        # 记录历史
        train_losses.append(avg_train_loss)
        val_losses.append(avg_val_loss)
        train_accuracies.append(train_accuracy)
        val_accuracies.append(val_accuracy)

        # 打印训练进度
        if epoch % 10 == 0 or epoch == num_epochs - 1:
            print(f"{epoch + 1:<6} {avg_train_loss:<10.4f} {avg_val_loss:<10.4f} "
                  f"{train_accuracy:<12.4f} {val_accuracy:<12.4f}")

        # 早停检查
        if avg_val_loss < best_val_loss:
            best_val_loss = avg_val_loss
            patience_counter = 0
            best_model_state = model.state_dict().copy()
        else:
            patience_counter += 1

        if patience_counter >= patience:
            print(f"\n早停触发! 验证损失在 {patience} 轮内没有改善")
            print(f"最佳验证损失: {best_val_loss:.4f} (第 {epoch + 1 - patience} 轮)")
            break

    # 加载最佳模型状态
    if best_model_state is not None:
        model.load_state_dict(best_model_state)
        print(f"✓ 已恢复到最佳模型状态")

    print(f"\n训练完成! 总共训练了 {epoch + 1} 轮")

    return model, {
        'train_losses': train_losses,
        'val_losses': val_losses,
        'train_accuracies': train_accuracies,
        'val_accuracies': val_accuracies
    }


def visualize_training_history(history):
    """
    可视化训练过程
    学习如何监控和分析模型训练过程
    """
    print("\n" + "=" * 50)
    print("步骤2.4: 训练过程可视化")
    print("=" * 50)

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))

    # 损失曲线
    ax1.plot(history['train_losses'], label='训练损失', color='blue')
    ax1.plot(history['val_losses'], label='验证损失', color='red')
    ax1.set_title('模型损失曲线')
    ax1.set_xlabel('训练轮次')
    ax1.set_ylabel('损失值')
    ax1.legend()
    ax1.grid(True, alpha=0.3)

    # 准确率曲线
    ax2.plot(history['train_accuracies'], label='训练准确率', color='blue')
    ax2.plot(history['val_accuracies'], label='验证准确率', color='red')
    ax2.set_title('模型准确率曲线')
    ax2.set_xlabel('训练轮次')
    ax2.set_ylabel('准确率')
    ax2.legend()
    ax2.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.show()

    # 分析训练过程
    final_train_acc = history['train_accuracies'][-1]
    final_val_acc = history['val_accuracies'][-1]
    acc_gap = abs(final_train_acc - final_val_acc)

    print(f"训练结果分析:")
    print(f"  最终训练准确率: {final_train_acc:.4f}")
    print(f"  最终验证准确率: {final_val_acc:.4f}")
    print(f"  准确率差距: {acc_gap:.4f}")

    if acc_gap > 0.05:
        print(f"  ⚠️  可能存在过拟合 (训练和验证准确率差距较大)")
    else:
        print(f"  ✓ 模型训练良好 (训练和验证准确率接近)")


def evaluate_model(model, test_loader, device=None):
    """
    在测试集上评估模型
    学习如何客观评估模型的最终性能
    """
    print("\n" + "=" * 50)
    print("步骤2.5: 测试集评估")
    print("=" * 50)

    if device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model.eval()
    all_predictions = []
    all_labels = []
    all_probabilities = []

    print("在测试集上进行预测...")

    with torch.no_grad():
        for batch_X, batch_y in test_loader:
            batch_X, batch_y = batch_X.to(device), batch_y.to(device)

            # 获取模型输出
            outputs = model(batch_X)
            probabilities = F.softmax(outputs, dim=1)

            # 获取预测结果
            _, predicted = torch.max(outputs, 1)

            # 收集结果
            all_predictions.extend(predicted.cpu().numpy())
            all_labels.extend(batch_y.cpu().numpy())
            all_probabilities.extend(probabilities.cpu().numpy())

    # 转换为numpy数组
    y_true = np.array(all_labels)
    y_pred = np.array(all_predictions)
    y_prob = np.array(all_probabilities)

    # 计算评估指标
    accuracy = accuracy_score(y_true, y_pred)
    precision = precision_score(y_true, y_pred)
    recall = recall_score(y_true, y_pred)
    f1 = f1_score(y_true, y_pred)

    print(f"测试集性能评估:")
    print(f"  准确率 (Accuracy): {accuracy:.4f}")
    print(f"  精确率 (Precision): {precision:.4f}")
    print(f"  召回率 (Recall): {recall:.4f}")
    print(f"  F1分数: {f1:.4f}")

    # 解释评估指标
    print(f"\n指标含义:")
    print(f"  准确率: 所有预测中正确的比例")
    print(f"  精确率: 预测为可饮用中真正可饮用的比例")
    print(f"  召回率: 所有可饮用水中被正确识别的比例")
    print(f"  F1分数: 精确率和召回率的调和平均")

    return {
        'accuracy': accuracy,
        'precision': precision,
        'recall': recall,
        'f1': f1,
        'predictions': y_pred,
        'probabilities': y_prob,
        'true_labels': y_true
    }


def save_trained_model(model, history, metrics, output_dir="trained_model"):
    """
    保存训练好的模型
    便于后续步骤使用和部署
    """
    import os

    print("\n" + "=" * 50)
    print("步骤2.6: 保存训练模型")
    print("=" * 50)

    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)

    # 保存模型状态（推荐方式，更安全）
    torch.save(model.state_dict(), f"{output_dir}/model_state.pt")
    print(f"✓ 模型状态已保存为: {output_dir}/model_state.pt")

    # 为兼容性也保存完整模型（在新版本PyTorch中可能有安全警告）
    try:
        torch.save(model, f"{output_dir}/complete_model.pt")
        print(f"✓ 完整模型已保存为: {output_dir}/complete_model.pt")
    except Exception as e:
        print(f"⚠️  完整模型保存失败（这是正常的）: {e}")

    # 保存训练历史和评估指标（处理numpy数组）
    # 注意：包含numpy数组的字典在PyTorch 2.6+中需要特殊处理
    try:
        # 尝试保存，如果失败则转换格式
        torch.save(history, f"{output_dir}/training_history.pt")
        torch.save(metrics, f"{output_dir}/evaluation_metrics.pt")
        print("✓ 训练历史和评估指标保存成功")
    except Exception as e:
        print(f"⚠️  直接保存失败，尝试格式转换: {e}")

        # 转换包含numpy数组的数据为纯Python类型
        history_converted = {}
        for key, value in history.items():
            if isinstance(value, list):
                history_converted[key] = [float(x) for x in value]
            else:
                history_converted[key] = value

        metrics_converted = {}
        for key, value in metrics.items():
            if hasattr(value, 'tolist'):  # numpy数组
                metrics_converted[key] = value.tolist()
            else:
                metrics_converted[key] = float(value) if isinstance(value, (np.float32, np.float64)) else value

        torch.save(history_converted, f"{output_dir}/training_history.pt")
        torch.save(metrics_converted, f"{output_dir}/evaluation_metrics.pt")
        print("✓ 格式转换后保存成功")

    # 保存模型架构信息（方便后续加载）
    model_info = {
        'input_size': model.fc1.in_features,
        'hidden_size': model.fc1.out_features,
        'dropout_rate': model.dropout1.p if hasattr(model.dropout1, 'p') else 0.3,
        'model_class': 'WaterQualityClassifier'
    }
    torch.save(model_info, f"{output_dir}/model_info.pt")

    print(f"✓ 模型相关文件已保存到 {output_dir}/ 目录:")
    print(f"  - model_state.pt: 模型参数状态 (推荐加载方式)")
    print(f"  - complete_model.pt: 完整模型 (兼容性)")
    print(f"  - training_history.pt: 训练历史记录")
    print(f"  - evaluation_metrics.pt: 评估指标")
    print(f"  - model_info.pt: 模型架构信息")
    print(f"\n💡 在后续步骤中，优先使用 model_state.pt 加载模型")


def main():
    """
    主函数 - 执行完整的监督分类任务构建流程
    这是机器学习项目的核心步骤
    """
    print("🎯 水质安全预测项目 - 监督分类任务构建")
    print("📚 学习目标: 掌握监督学习分类任务的完整流程")
    print("包括: 数据划分、模型设计、训练过程、性能评估")

    try:
        # 步骤1: 加载预处理数据
        X, y, scaler = load_preprocessed_data()
        if X is None:
            return

        # 步骤2: 划分数据集
        X_train, X_val, X_test, y_train, y_val, y_test = split_dataset(X, y)

        # 步骤3: 创建数据加载器
        train_loader, val_loader, test_loader = create_data_loaders(
            X_train, y_train, X_val, y_val, X_test, y_test
        )

        # 步骤4: 创建模型
        input_size = X.shape[1]  # 特征数量
        model = WaterQualityClassifier(input_size=input_size)

        print(f"\n模型架构:")
        print(model)

        # 计算模型参数数量
        total_params = sum(p.numel() for p in model.parameters())
        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
        print(f"模型参数总数: {total_params}")
        print(f"可训练参数: {trainable_params}")

        # 步骤5: 训练模型
        trained_model, history = train_model(model, train_loader, val_loader)

        # 步骤6: 可视化训练过程
        visualize_training_history(history)

        # 步骤7: 评估模型
        metrics = evaluate_model(trained_model, test_loader)

        # 步骤8: 保存模型
        save_trained_model(trained_model, history, metrics)

        print("\n" + "🎉" * 20)
        print("监督分类任务构建完成!")
        print("下一步: 进行步骤3 - 比较不同算法性能")
        print("🎉" * 20)

    except Exception as e:
        print(f"❌ 训练过程中发生错误: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()

"""
📝 学习总结:
通过这个步骤，你学会了:
1. 如何将实际问题转化为监督学习分类任务
2. 正确的数据集划分策略 (训练/验证/测试)
3. 使用PyTorch设计和实现神经网络模型
4. 监督学习的训练过程 (前向传播、反向传播、参数更新)
5. 如何监控训练过程并防止过拟合
6. 使用标准指标评估分类模型性能

这些技能是深度学习和机器学习的核心!
"""