main/hyperpara_optim.py

import os
import torch
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import seaborn as sns
from torch.utils.data import DataLoader
import random
import optuna
import gc # [新增] 用于垃圾回收
import warnings # [新增] 用于屏蔽警告

from Ramandataset import RamanDataset
# 导入您提供的 Raman_Task 中的必要组件
from Raman_Task import (
    load_real_data, 
    load_class_names, 
    stratified_split_with_minimum_samples,
    augment_minority_classes,
    load_mae_model_for_classification,
    train_predictor,
    unique
)

def mm_to_inches(mm):
    return mm / 25.4

def run_hyperparameter_optimization(
    num_trials=10,
    epochs_per_trial=50,
    resume_existing=False,
    study_suffix=None,
):
    """
    执行超参数随机搜索，并生成报告和可视化图表。
    """
    
    # 1. 定义超参数搜索空间
    # (这部分逻辑移动到了 objective 内部，这里仅作说明)
    
    # 2. 数据准备 (只需加载一次原始数据，避免重复IO)
    base_path = '/home/lion/Desktop/splendid_lion/project/2025/multi_source/Raman/support_file/data/'
    
    data_path = os.path.join(base_path, 'rruff/classifier_0_3500_spectra.npy') 
    labels_path = os.path.join(base_path, 'rruff/classifier_0_3500_labels.npy') 
    wavenumbers_path = os.path.join(base_path, 'rruff/classifier_0_3500_wavenumbers.npy') 

    print("Loading data for optimization...")
    spectra, labels, wavenumbers = load_real_data(data_path, labels_path=labels_path, wavenumbers_path=wavenumbers_path, normalize=True)
    num_classes = len(unique(labels))
    input_length = spectra.shape[1]
    
    # 固定数据划分，确保比较的公平性
    from sklearn.preprocessing import LabelEncoder
    y = np.array(labels)
    le = LabelEncoder()
    y_encoded = le.fit_transform(y)
    
    X_train, X_val, X_test, y_train, y_val, y_test = stratified_split_with_minimum_samples(
        spectra, y_encoded, test_size=0.15, val_size=0.15, min_samples_per_class=1, random_state=42
    )
    
    # 数据增强
    X_train_augmented, y_train_augmented = augment_minority_classes(
        X_train, y_train, min_samples=30, target_samples=159
    )
    
    # 3. 定义 Optuna Objective
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"\n🚀 Starting Hyperparameter Optimization with Optuna ({num_trials} trials)...")
    print(f"Device: {device}")
    
    # 用于存储所有试验结果以供分析
    results_list = []

    def objective(trial):
        # [新增] 1. 每个 Trial 开始前，强制清理之前的显存碎片
        gc.collect()
        torch.cuda.empty_cache()

        # 定义搜索空间
        lr = trial.suggest_categorical('lr', [1e-3, 5e-4, 1e-4, 5e-5, 1e-5])
        batch_size = trial.suggest_categorical('batch_size', [16, 32])
        weight_decay = trial.suggest_categorical('weight_decay', [1e-3, 5e-4, 1e-4, 5e-5, 1e-5])
        mask_ratio = trial.suggest_categorical('mask_ratio', [0.25, 0.5, 0.75])

        print(f"\n--- Trial {trial.number} ---")
        print(f"Params: lr={lr}, batch_size={batch_size}, weight_decay={weight_decay}, mask_ratio={mask_ratio}")

        # 初始化变量，防止 finally 块报错
        classifier = None
        mae_model = None
        train_loader = None
        val_loader = None
        test_loader = None

        try:
            # 重新构建 DataLoaders
            # [注意] num_workers=0 是关键，AMD 显卡上多进程 DataLoader 容易导致 Segfault
            train_dataset = RamanDataset(X_train_augmented, None, labels=y_train_augmented, transform=None, is_train=True)
            train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=True, num_workers=0)
            
            val_dataset = RamanDataset(X_val, None, labels=y_val, transform=None, is_train=True)
            val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=0)
            
            test_dataset = RamanDataset(X_test, None, labels=y_test, transform=None, is_train=True)
            test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=0)
            
            # 加载模型
            para = {"embedding_dim": 512, "num_heads": 16, "num_layers": 12, "patch_num": 100}
            pretrained_path = (
                f"/home/lion/Desktop/splendid_lion/project/2025/multi_source/Raman/support_file/model/"
                f"0.0001_{mask_ratio}_{para['embedding_dim']}_{para['num_heads']}_{para['num_layers']}_{para['patch_num']}/"
                f"Fine_tuned_baseonALL.pth"
            )
            
            if not os.path.exists(pretrained_path):
                print(f"⚠️ Trial {trial.number} pruned: pretrained model not found at {pretrained_path}")
                raise optuna.exceptions.TrialPruned()

            classifier, encoder, decoder, mae_model = load_mae_model_for_classification(
                pretrained_path, input_length, para['patch_num'], para['embedding_dim'],
                para['num_layers'], para['num_heads'], num_classes, device
            )
            
            save_dir = os.path.join(base_path, f"optimization_results/trial_{trial.number}")
            os.makedirs(save_dir, exist_ok=True)
            
            # 训练模型
            # 注意：如果 Raman_Task.py 中已经实现了梯度裁剪和 .contiguous()，这里会更稳定
            trained_model, _ = train_predictor(
                classifier=classifier,
                mae_model=mae_model,
                train_loader=train_loader,
                val_loader=val_loader,
                test_loader=test_loader,
                device=device,
                epochs=epochs_per_trial,
                lr=lr,
                weight_decay=weight_decay,
                patience=10,
                save_dir=save_dir,
                model_name=f"trial_{trial.number}",
                freeze_encoder=False
            )
            
            # 手动计算验证集准确率
            classifier.eval()
            correct = 0
            total = 0
            with torch.no_grad():
                for inputs, _, labels in val_loader:
                    inputs, labels = inputs.to(device), labels.to(device)
                    outputs, _ = classifier(inputs)
                    _, predicted = torch.max(outputs.data, 1)
                    total += labels.size(0)
                    correct += (predicted == labels.squeeze()).sum().item()
            
            val_acc = correct / total
            
            # 记录结果
            result_record = {
                'trial_id': trial.number,
                'val_accuracy': val_acc,
                'lr': lr,
                'batch_size': batch_size,
                'weight_decay': weight_decay,
                'mask_ratio': mask_ratio
            }
            results_list.append(result_record)

            # 立即保存
            csv_file = "optuna_results.csv"
            df_current = pd.DataFrame([result_record])
            if not os.path.exists(csv_file):
                df_current.to_csv(csv_file, index=False, mode='w')
            else:
                df_current.to_csv(csv_file, index=False, mode='a', header=False)
            print(f"✅ Trial {trial.number} Finished. Acc: {val_acc:.4f}")
            
            return val_acc

        # [新增] 2. 捕获运行时错误 (OOM 或 NaN)，防止整个程序崩溃
        except RuntimeError as e:
            error_msg = str(e)
            if "out of memory" in error_msg:
                print(f"⚠️ Trial {trial.number} Pruned due to CUDA OOM.")
                # 抛出 TrialPruned 异常，Optuna 会标记此试验为被剪枝，而不是失败
                raise optuna.exceptions.TrialPruned()
            elif "nan" in error_msg.lower(): # 捕获 Loss 变成 NaN 的情况
                print(f"⚠️ Trial {trial.number} Pruned due to NaN loss.")
                raise optuna.exceptions.TrialPruned()
            else:
                print(f"❌ Trial {trial.number} Failed with RuntimeError: {e}")
                return 0.0 # 或者抛出异常
        
        except Exception as e:
            print(f"❌ Trial {trial.number} Failed with unknown error: {e}")
            return 0.0
            
        finally:
            # [新增] 3. 暴力清理，确保为下一个 Trial 腾出显存
            if classifier is not None:
                del classifier
            if mae_model is not None:
                del mae_model
            # 清理 DataLoaders (有时它们会持有显存引用)
            del train_loader, val_loader, test_loader
            
            gc.collect()
            torch.cuda.empty_cache()
            print(f"🧹 Trial {trial.number} Cleanup Done.")

    # 4. 创建 Study 并开始优化
    if study_suffix is None:
        study_suffix = pd.Timestamp.now().strftime("%Y%m%d_%H%M%S")
    if resume_existing:
        study_name = "raman_optimization"
    else:
        study_name = f"raman_optimization_fresh_{study_suffix}"
    
    # 建议使用 SQLite 存储，这样即使程序崩溃，之前的 Trial 记录还在
    storage_url = "sqlite:///{}.db".format(study_name)
    
    study = optuna.create_study(
        direction="maximize", 
        study_name=study_name,
        storage=storage_url, # 启用持久化存储
        load_if_exists=resume_existing  # False 时从头开始，不加载历史结果
    )
    
    # [新增] n_jobs=1 是必须的，防止多进程导致 ROCm 崩溃
    # gc_after_trial=True 让 Optuna 协助清理
    study.optimize(objective, n_trials=num_trials, n_jobs=1, gc_after_trial=True)

    # 5. 生成报告和可视化
    print("\n" + "="*50)
    print("HYPERPARAMETER OPTIMIZATION REPORT (Optuna)")
    print("="*50)
    
    if len(study.trials) > 0:
        best_trial = study.best_trial
        print(f"Total Trials: {len(study.trials)}")
        print(f"Best Validation Accuracy: {best_trial.value:.4f}")
        print("\nBest Hyperparameters:")
        for key, value in best_trial.params.items():
            print(f"   {key}: {value}")
    else:
        print("No successful trials.")
    print("="*50 + "\n")
    
    all_trials_data = []
    print(f"Total trials in DB: {len(study.trials)}")
    for t in study.trials:
        # 过滤条件：必须完成，且准确率大于 0 (排除 return 0.0 的失败情况)
        if t.state == optuna.trial.TrialState.COMPLETE and t.value is not None and t.value > 0.0001:
            all_trials_data.append({
                'trial_id': t.number,
                'val_accuracy': t.value,
                'lr': t.params.get('lr'),
                'batch_size': t.params.get('batch_size'),
                'weight_decay': t.params.get('weight_decay'),
                'mask_ratio': t.params.get('mask_ratio')
            })
            
    if all_trials_data:
        df_results = pd.DataFrame(all_trials_data)
        print(f"Visualizing {len(df_results)} total trials from database.") # 打印总数确认
        df_results.to_csv("optuna_results_final.csv", index=False)
        visualize_hyperparameters_full_width(df_results)
    else:
        print("No successful trials found in study.")

def visualize_hyperparameters_full_width(df):
    """
    绘制通栏大小 (170mm 宽度) 的超参数分析图。
    """
    # 保持原有绘图逻辑不变
    plt.rcParams.update({
        'font.size': 7,
        'axes.labelsize': 8,
        'xtick.labelsize': 7,
        'ytick.labelsize': 7,
        'legend.fontsize': 7,
        'font.family': 'sans-serif',
        'lines.linewidth': 1.0,
        'axes.linewidth': 0.8
    })
    
    fig, axes = plt.subplots(2, 2, figsize=(mm_to_inches(170), mm_to_inches(120)), constrained_layout=True)
    axes = axes.flatten()
    
    # 定义明显的中位线样式：加粗并使用深红色
    median_props = dict(linewidth=1.5, color='firebrick')

    # --- 子图 1: Learning Rate ---
    # [修改] 添加 medianprops 以突出中位线，即使与箱体边缘重合也能看见
    sns.boxplot(x='lr', y='val_accuracy', data=df, ax=axes[0], palette="Blues", linewidth=0.8, showfliers=False, medianprops=median_props)
    sns.stripplot(x='lr', y='val_accuracy', data=df, ax=axes[0], color='darkblue', alpha=0.6, jitter=0.1, size=4)
    axes[0].set_xlabel("Learning Rate")
    axes[0].set_ylabel("Validation Accuracy")
    axes[0].set_title("(a) Impact of Learning Rate")
    axes[0].grid(axis='y', linestyle='--', alpha=0.5)

    # --- 子图 2: Batch Size ---
    sns.boxplot(x='batch_size', y='val_accuracy', data=df, ax=axes[1], palette="Greens", linewidth=0.8, showfliers=False, medianprops=median_props)
    sns.stripplot(x='batch_size', y='val_accuracy', data=df, ax=axes[1], color='darkgreen', alpha=0.6, jitter=0.1, size=4)
    axes[1].set_xlabel("Batch Size")
    axes[1].set_ylabel("") 
    axes[1].set_title("(b) Impact of Batch Size")
    axes[1].grid(axis='y', linestyle='--', alpha=0.5)

    # --- 子图 3: Weight Decay ---
    sns.boxplot(x='weight_decay', y='val_accuracy', data=df, ax=axes[2], palette="Oranges", linewidth=0.8, showfliers=False, medianprops=median_props)
    sns.stripplot(x='weight_decay', y='val_accuracy', data=df, ax=axes[2], color='darkred', alpha=0.6, jitter=0.1, size=4)
    axes[2].set_xlabel("Weight Decay")
    axes[2].set_ylabel("")
    axes[2].set_title("(c) Impact of Weight Decay")
    axes[2].grid(axis='y', linestyle='--', alpha=0.5)

    # --- 子图 4: Mask Ratio ---
    sns.boxplot(x='mask_ratio', y='val_accuracy', data=df, ax=axes[3], palette="Purples", linewidth=0.8, showfliers=False, medianprops=median_props)
    sns.stripplot(x='mask_ratio', y='val_accuracy', data=df, ax=axes[3], color='indigo', alpha=0.6, jitter=0.1, size=4)
    axes[3].set_xlabel("Mask Ratio")
    axes[3].set_ylabel("")
    axes[3].set_title("(d) Impact of Mask Ratio")
    axes[3].grid(axis='y', linestyle='--', alpha=0.5)

    for ax in axes:
        for spine in ax.spines.values():
            spine.set_visible(True)
            spine.set_edgecolor('black')
            spine.set_linewidth(0.8)

    save_path = "hyperparameter_optimization_analysis.png"
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    print(f"Visualization saved to {save_path}")

if __name__ == "__main__":
    # 启用 SQLite 存储后，您可以随时中断程序，下次运行会继续之前的进度
    run_hyperparameter_optimization(num_trials=25, epochs_per_trial=30, resume_existing=False)