train.py

import os
import numpy as np
import pandas as pd
import tensorflow as tf
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder, StandardScaler
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Dropout
from tensorflow.keras.utils import to_categorical
from tensorflow.keras.callbacks import EarlyStopping
from sklearn.metrics import confusion_matrix
import seaborn as sns
import matplotlib.pyplot as plt

# 配置参数
DATA_DIR = "output"  # 数据目录
SEQ_LENGTH = 402  # 时间序列长度（402个数据点）
TEST_SIZE = 0.5  # 测试集比例
RANDOM_STATE = 42  # 随机种子
BATCH_SIZE = 16  # 批次大小
EPOCHS = 100  # 最大训练轮数


def load_data(data_dir):
    """加载数据集并提取特征和标签"""
    features = []
    labels = []

    # 遍历数据目录
    for filename in os.listdir(data_dir):
        if filename.endswith(".csv"):
            # 从文件名提取标签（XXX部分）
            label = filename.split("-")[-2]  # 根据实际文件名结构调整索引

            # 读取CSV文件
            filepath = os.path.join(data_dir, filename)
            df = pd.read_csv(filepath, usecols=[0], skiprows=1, nrows=SEQ_LENGTH)

            # 转换为numpy数组并存储
            features.append(df.values.flatten())
            labels.append(label)

    return np.array(features), np.array(labels)


def preprocess_data(features, labels):
    """数据预处理流程"""
    # 标签编码
    le = LabelEncoder()
    encoded_labels = le.fit_transform(labels)
    onehot_labels = to_categorical(encoded_labels)

    # 特征标准化
    scaler = StandardScaler()
    scaled_features = scaler.fit_transform(features)

    return scaled_features, onehot_labels, le.classes_


def build_model(input_shape, num_classes):
    """构建ANN模型"""
    model = Sequential([
        Dense(512, activation='relu', input_shape=(input_shape,)),
        Dropout(0.3),
        Dense(256, activation='relu'),
        Dropout(0.3),
        Dense(128, activation='relu'),
        Dense(num_classes, activation='softmax')
    ])

    model.compile(optimizer='adam',
                  loss='categorical_crossentropy',
                  metrics=['accuracy'])
    return model


def main():
    # 1. 加载数据
    print("正在加载数据...")
    features, labels = load_data(DATA_DIR)

    # 2. 数据预处理
    print("\n数据预处理...")
    X, y, class_names = preprocess_data(features, labels)

    # 3. 划分数据集
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=TEST_SIZE, stratify=y, random_state=RANDOM_STATE
    )

    # 4. 构建模型
    print("\n构建模型...")
    model = build_model(X_train.shape[1], y.shape[1])
    model.summary()

    # 5. 训练配置
    early_stop = EarlyStopping(monitor='val_loss', patience=10, restore_best_weights=True)

    # 6. 训练模型
    print("\n开始训练...")
    history = model.fit(
        X_train, y_train,
        validation_split=0.2,
        epochs=EPOCHS,
        batch_size=BATCH_SIZE,
        callbacks=[early_stop],
        verbose=1
    )

    # 7. 模型评估
    print("\n模型评估：")
    test_loss, test_acc = model.evaluate(X_test, y_test, verbose=0)
    y_pred = model.predict(X_test)
    cm = confusion_matrix(np.argmax(y_test, axis=1), np.argmax(y_pred, axis=1))
    # plt.figure(figsize=(10, 7))
    # sns.heatmap(cm, annot=True, fmt='d', xticklabels=class_names, yticklabels=class_names)
    # plt.xlabel('Predicted')
    # plt.ylabel('True')
    # plt.title('Confusion Matrix')
    # plt.show()
    # print(f"测试集准确率：{test_acc:.4f}")
    # print(f"测试集损失值：{test_loss:.4f}")
    plt.figure(figsize=(10, 7))
    # 使用annot_kws调整矩阵内数字的字号，fontsize调整坐标轴标签字号
    sns.heatmap(cm, annot=True, fmt='d',
                annot_kws={"size": 16},  # 矩阵内数字字号
                xticklabels=class_names,
                yticklabels=class_names,
                cbar_kws={"shrink": 0.8})  # 可选的色条大小调整

    # 设置坐标轴标签字号
    plt.xlabel('Predicted', fontsize=14)
    plt.ylabel('True', fontsize=14)
    plt.title('Confusion Matrix', fontsize=16)

    # 调整坐标轴刻度标签字号
    plt.xticks(fontsize=12)
    plt.yticks(fontsize=12)

    plt.tight_layout()  # 自动调整布局
    plt.show()

    print(f"测试集准确率：{test_acc:.4f}")
    print(f"测试集损失值：{test_loss:.4f}")

    # 8. 保存模型
    model.summary()
    model.save("timeseries_classifier.h5")
    print("\n模型已保存为 timeseries_classifier.h5")
    y_pred = model.predict(X_test)
    cm = confusion_matrix(np.argmax(y_test, axis=1), np.argmax(y_pred, axis=1))
    sns.heatmap(cm, annot=True, fmt='d', xticklabels=class_names, yticklabels=class_names)


if __name__ == "__main__":
    main()