train_esi_function.py

import os
import glob
import traceback
import torch
import torch.nn as nn
from torch.cuda.amp import autocast, GradScaler  # 混合精度
from torch.utils.data import DataLoader, IterableDataset
import pandas as pd
import numpy as np
import json
import time
from datetime import datetime
from tqdm import tqdm
import matplotlib.pyplot as plt
from Enzoria.model.enzoria.enzoria_affinity_model import EnzoriaAffinityModel
import re
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score, roc_auc_score


PROJECT_NAME = "Enzoria"

class TrainingLogger:
    def __init__(self, result_dir="./result", run_name=None):
        # 生成唯一运行ID：时间戳 + 自定义名
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        if run_name:
            run_id = f"run_{run_name}_{timestamp}"
        else:
            run_id = f"run_{timestamp}"
        
        # 创建唯一运行目录（基目录下）
        self.run_dir = os.path.join(result_dir, run_id)
        self.log_dir = os.path.join(self.run_dir, "logs")
        self.checkpoint_dir = os.path.join(self.run_dir, "checkpoints")
        self.plot_dir = os.path.join(self.run_dir, "plots")
        
        # 处理目录冲突：如果存在，添加增量编号
        counter = 1
        base_run_dir = self.run_dir
        while os.path.exists(self.run_dir):
            self.run_dir = f"{base_run_dir}_{counter}"
            self.log_dir = os.path.join(self.run_dir, "logs")
            self.checkpoint_dir = os.path.join(self.run_dir, "checkpoints")
            self.plot_dir = os.path.join(self.run_dir, "plots")
            counter += 1
        
        # 创建目录
        os.makedirs(self.log_dir, exist_ok=True)
        os.makedirs(self.checkpoint_dir, exist_ok=True)
        os.makedirs(self.plot_dir, exist_ok=True)
        
        print(f"📂 新运行目录创建: {self.run_dir}")
        
        # 初始化历史记录（同原）
        self.batch_history = {
            'epoch': [],
            'batch': [],
            'loss': [],
            'lr': [],
            'timestamp': []  # 新增：每个log的时间戳
        }
        
        self.epoch_history = {
            'epoch': [],
            'metrics': [],  # 存储指标字典
            'lr': [],
            'timestamp': []  # 新增
        }
    
    def log_batch_loss(self, epoch, batch, loss, lr):
        current_time = time.time()
        self.batch_history['epoch'].append(epoch)
        self.batch_history['batch'].append(batch)
        self.batch_history['loss'].append(loss)
        self.batch_history['lr'].append(lr)
        self.batch_history['timestamp'].append(current_time)
        
        # 保存JSON（完整历史）
        with open(os.path.join(self.log_dir, "batch_history.json"), 'w') as f:
            json.dump(self.batch_history, f, indent=2)
        
        # 新增：追加到CSV（便于分析）
        batch_df = pd.DataFrame([{
            'epoch': epoch,
            'batch': batch,
            'loss': loss,
            'lr': lr,
            'timestamp': current_time
        }])
        batch_csv_path = os.path.join(self.log_dir, "batch_history.csv")
        batch_df.to_csv(batch_csv_path, mode='a', header=not os.path.exists(batch_csv_path), index=False)
    
    def log_epoch(self, epoch, metrics, lr, prefix=""):
        current_time = time.time()
        metrics_with_prefix = {f"{prefix}{k}": v for k, v in metrics.items()}
        
        idx = next((i for i, e in enumerate(self.epoch_history['epoch']) if e == epoch), None)
        
        if idx is not None:
            # 更新现有条目
            self.epoch_history['metrics'][idx].update(metrics_with_prefix)
            self.epoch_history['timestamp'][idx] = current_time
            # lr 如果需要更新，但假设 train 先设置
        else:
            # 添加新条目
            self.epoch_history['epoch'].append(epoch)
            self.epoch_history['metrics'].append(metrics_with_prefix)
            self.epoch_history['lr'].append(lr)
            self.epoch_history['timestamp'].append(current_time)
        
        # 保存JSON
        with open(os.path.join(self.log_dir, "epoch_history.json"), 'w') as f:
            json.dump(self.epoch_history, f, indent=2)
        
        # 新增：追加到CSV（每个log作为一个row）
        epoch_df = pd.DataFrame([{
            'epoch': epoch,
            **metrics_with_prefix,  # 展开metrics dict
            'lr': lr,
            'timestamp': current_time
        }])
        epoch_csv_path = os.path.join(self.log_dir, "epoch_history.csv")
        epoch_df.to_csv(epoch_csv_path, mode='a', header=not os.path.exists(epoch_csv_path), index=False)
    
    def plot_loss_curves(self):
        """绘制所有指标的曲线图（同原，但路径用self.plot_dir）"""
        plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans', 'Arial Unicode MS']
        plt.rcParams['axes.unicode_minus'] = False

        epochs = self.epoch_history['epoch']
        metrics_list = self.epoch_history['metrics']
        
        if len(epochs) < 2:
            print("⚠️ 需要至少2个epoch才能绘制指标曲线")
            return
        
        # 提取各个指标
        train_loss = [m.get('loss', np.nan) for m in metrics_list]
        val_loss = [m.get('val_loss', np.nan) for m in metrics_list]
        
        train_acc = [m.get('acc', np.nan) for m in metrics_list]
        val_acc = [m.get('val_acc', np.nan) for m in metrics_list]
        
        train_f1 = [m.get('f1', np.nan) for m in metrics_list]
        val_f1 = [m.get('val_f1', np.nan) for m in metrics_list]
        
        train_auc = [m.get('auc', np.nan) for m in metrics_list]
        val_auc = [m.get('val_auc', np.nan) for m in metrics_list]
        
        # 创建子图
        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 10))
        
        # Loss
        ax1.plot(epochs, train_loss, 'b-', label='Train Loss', linewidth=2)
        ax1.plot(epochs, val_loss, 'r--', label='Val Loss', linewidth=2)
        ax1.set_xlabel('Epoch')
        ax1.set_ylabel('Loss')
        ax1.set_title('Loss (Train & Val)')
        ax1.legend()
        ax1.grid(True, alpha=0.3)
        
        # Accuracy
        ax2.plot(epochs, train_acc, 'b-', label='Train Acc', linewidth=2)
        ax2.plot(epochs, val_acc, 'r--', label='Val Acc', linewidth=2)
        ax2.set_xlabel('Epoch')
        ax2.set_ylabel('Accuracy')
        ax2.set_title('Accuracy (Train & Val)')
        ax2.legend()
        ax2.grid(True, alpha=0.3)
        
        # F1
        ax3.plot(epochs, train_f1, 'b-', label='Train F1', linewidth=2)
        ax3.plot(epochs, val_f1, 'r--', label='Val F1', linewidth=2)
        ax3.set_xlabel('Epoch')
        ax3.set_ylabel('F1 Score')
        ax3.set_title('F1 Score (Train & Val)')
        ax3.legend()
        ax3.grid(True, alpha=0.3)
        
        # AUC
        ax4.plot(epochs, train_auc, 'b-', label='Train AUC', linewidth=2)
        ax4.plot(epochs, val_auc, 'r--', label='Val AUC', linewidth=2)
        ax4.set_xlabel('Epoch')
        ax4.set_ylabel('AUC')
        ax4.set_title('AUC (Train & Val)')
        ax4.legend()
        ax4.grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.savefig(os.path.join(self.plot_dir, 'metrics_curves.png'), dpi=300, bbox_inches='tight')
        plt.close()
        
        print(f"📈 指标曲线已保存: {os.path.join(self.plot_dir, 'metrics_curves.png')}")
    
    def save_training_summary(self, config, total_time):
        """保存训练总结（同原，但路径用self.log_dir，并添加run_dir）"""
        summary = {
            'run_dir': self.run_dir,  # 新增：记录运行目录
            'config': config,
            'total_training_time': total_time,
            'total_epochs': len(self.epoch_history['epoch']),
            'final_metrics': self.epoch_history['metrics'][-1] if self.epoch_history['metrics'] else {},
            'best_metrics': {
                'best_val_f1': max([m.get('val_f1', 0) for m in self.epoch_history['metrics']]) if self.epoch_history['metrics'] else None,
                'best_val_acc': max([m.get('val_acc', 0) for m in self.epoch_history['metrics']]) if self.epoch_history['metrics'] else None,
                'best_val_auc': max([m.get('val_auc', 0) for m in self.epoch_history['metrics']]) if self.epoch_history['metrics'] else None,
            },
            'test_metrics': {
                'test_f1': self.epoch_history['metrics'][-1].get('test_f1', None),
                'test_acc': self.epoch_history['metrics'][-1].get('test_acc', None),
                'test_auc': self.epoch_history['metrics'][-1].get('test_auc', None),
            }
        }
        
        with open(os.path.join(self.log_dir, "training_summary.json"), 'w') as f:
            json.dump(summary, f, indent=2)
        
        print(f"📄 训练总结已保存: {os.path.join(self.log_dir, 'training_summary.json')}")

class StreamingPDBBindDataset(IterableDataset):
    """
    流式数据集，逐块加载大CSV文件，避免全载入内存。
    """
    def __init__(self, file_paths, chunksize=500, shuffle=False):  # 减小chunksize以降低内存
        self.file_paths = file_paths
        self.chunksize = chunksize
        self.shuffle = shuffle  # 注意：全shuffle需全载；此处近似chunk内shuffle

    def __iter__(self):
        for file_path in self.file_paths:
            reader = pd.read_csv(file_path, chunksize=self.chunksize, dtype=str)
            for chunk in reader:
                chunk = chunk.dropna(subset=["sequence", "foldseek_seq", "text", "label"])
                if self.shuffle:
                    chunk = chunk.sample(frac=1).reset_index(drop=True)
                for _, row in chunk.iterrows():
                    yield {
                        "sequence": row["sequence"],
                        "structure": row["foldseek_seq"],
                        "text": row["text"],
                        "label": torch.tensor(float(row["label"]), dtype=torch.float)  # BCE需float
                    }

# === 主训练函数 ===
def train_function_model(  # 修改函数名为二分类任务
    base_dir="dataset/ESIbank/processed_splits",
    result_dir="./result",
    batch_size=8,  # 增加到8，利用更大内存加速
    accumulation_steps=1,  # 相应减少，保持有效batch=8
    num_epochs=5,
    lr=2e-5,
    device="cuda"
):
    # 初始化日志
    logger = TrainingLogger(result_dir)
    
    # 配置（添加max_lengths减小序列长）
    config = {
        "protein_config": "./Enzoria/weights/Enzoria_650M/esm2_t33_650M_UR50D",
        "text_config": "./Enzoria/weights/Enzoria_650M/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext",
        "structure_config": "./Enzoria/weights/Enzoria_650M/foldseek_t30_150M",
        "load_protein_pretrained": False,
        "load_text_pretrained": False,
        "from_checkpoint": "./Enzoria/weights/Enzoria_650M/Enzoria_650M.pt",
        "optimizer_kwargs": {
            "class": "AdamW",
            "betas": (0.9, 0.999),
            "weight_decay": 0.01,
        },
        "lr_scheduler_kwargs": {
            "class": "CosineAnnealingWarmRestarts",
            "init_lr": 5e-5,
            "T_0": 5,
            "T_mult": 2,
            "eta_min": 1e-7,
        },
    }

    # 保存配置
    with open(os.path.join(logger.log_dir, "training_config.json"), 'w') as f:
        json.dump(config, f, indent=2)
    print(f"{PROJECT_NAME}: starting ESIBank binary classification finetuning...")
    
    print("🚀 开始ESIBank功能二分类模型微调...")
    print(f"📁 结果保存至: {result_dir}")
    print(f"📊 训练配置: {json.dumps(config, indent=2)}")
    
    model = EnzoriaAffinityModel(**config).to(device)
    
    # 加载检查点
    if "from_checkpoint" in config and os.path.exists(config["from_checkpoint"]):
        print(f"📥 加载预训练检查点: {config['from_checkpoint']}")
        checkpoint = torch.load(config["from_checkpoint"], map_location=device)
        model.load_state_dict(checkpoint, strict=False)
    else:
        print("⚠️ 未找到 from_checkpoint 或路径不存在")

    # =========================================================
    # 控制训练策略（可选 LoRA、部分解冻、embedding 开放）
    # =========================================================
    freeze_encoder = config.get("freeze_encoder", False)
    open_last_layers = config.get("open_last_layers", 4)   # 开放最后N层
    tune_embedding = config.get("tune_embedding", True)    # 是否开放embedding层

    trainable_params = []
  
        
    if freeze_encoder:
        print("🔒 冻结编码器参数，只训练预测头...")
        for name, param in model.named_parameters():
            if "affinity_head" not in name:
                param.requires_grad = False
            else:
                param.requires_grad = True
                trainable_params.append(name)

    else:
        print(f"🧩 开放 ESM 最后 {open_last_layers} 层 & 预测头...")

        # ESM 模型的层范围（从你的参数列表看是11-32）
        esm_start_layer = 11  # 第一层的编号
        esm_total_layers = 22  # 总共22层 (11-32)
        start_layer = esm_start_layer + (esm_total_layers - open_last_layers)

        for name, param in model.named_parameters():
            param.requires_grad = False

            # 1️⃣ 预测头
            if "affinity_head" in name:
                param.requires_grad = True
                trainable_params.append(name)

            # 2️⃣ ESM 编码器最后几层
            elif "protein_encoder.model.esm.encoder.layer" in name:
                import re
                layer_match = re.search(r'layer\.(\d+)', name)
                if layer_match:
                    layer_num = int(layer_match.group(1))
                    if layer_num >= start_layer:
                        param.requires_grad = True
                        trainable_params.append(name)

            # 3️⃣ ESM 输出投影层
            elif "protein_encoder.out" in name:
                param.requires_grad = True
                trainable_params.append(name)

            # 4️⃣ Embedding 层（可选）
            elif tune_embedding and any(x in name for x in ["embed_tokens", "embeddings"]):
                param.requires_grad = True
                trainable_params.append(name)

            # 5️⃣ LayerNorm 层（可选，通常建议开放）
            elif "LayerNorm" in name and "protein_encoder" in name:
                param.requires_grad = True
                trainable_params.append(name)

    print(f"🎯 可训练参数 ({len(trainable_params)}):")
    # 按模块分组显示
    modules = {}
    for pname in trainable_params:
        module = pname.split('.')[0] if '.' in pname else pname
        if module not in modules:
            modules[module] = []
        modules[module].append(pname)

    for module, params in modules.items():
        print(f"📦 {module}: {len(params)} 个参数")
        for pname in params[:3]:  # 每个模块显示前3个
            print(f"   - {pname}")
        if len(params) > 3:
            print(f"   ... 还有 {len(params) - 3} 个参数")
    
    model.train()
    model.training_mode = "finetune"  # 确保模式为 finetune

    # 2️⃣ 优化器 + scheduler + 损失函数
    optim_dict = model.configure_optimizers()
    optimizer = optim_dict["optimizer"]
    scheduler = optim_dict["lr_scheduler"]["scheduler"]
    criterion = nn.BCEWithLogitsLoss()  # 使用sigmoid的二分类损失
    scaler = GradScaler()  # AMP
    # 3️⃣ 数据加载
    split_types = ["all_split", "enzyme_split", "random_split", "reaction_split"]
    base_dir = "dataset/ESIbank/processed_splits"
    result_dir = "./result"

    for split_type in split_types:
        split_dir = os.path.join(base_dir, split_type)
        split_result_dir = os.path.join(result_dir, split_type)
        os.makedirs(split_result_dir, exist_ok=True)

        print(f"\n===== 🧩 当前分割类型: {split_type} =====")

        def get_file_paths(data_type: str):
            """辅助函数：获取某个数据类型（training/val/testing）的 CSV 文件路径列表"""
            files = sorted(glob.glob(os.path.join(split_dir, f"{data_type}_datas_*.csv")))
            if not files:
                print(f"⚠️ 未找到 {data_type} 数据文件")
            return files

        # === 获取三种数据集的文件路径 ===
        train_files = get_file_paths("training")
        val_files = get_file_paths("val")
        test_files = get_file_paths("testing")

        if not train_files:
            print(f"❌ {split_type} 缺少训练集，跳过。")
            continue

        # === 构建 Dataset & DataLoader（流式加载） ===
        train_dataset = StreamingPDBBindDataset(train_files, chunksize=500, shuffle=True)  # shuffle为chunk内近似
        val_dataset = StreamingPDBBindDataset(val_files, chunksize=500) if val_files else None
        test_dataset = StreamingPDBBindDataset(test_files, chunksize=500) if test_files else None

        train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=False, num_workers=4)  # 增加num_workers加速加载
        val_dataloader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4) if val_dataset else None
        test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=4) if test_dataset else None

        # === 初始化日志器 ===
        split_logger = TrainingLogger(split_result_dir)

        print(f"✅ {split_type} 数据加载完成 (流式模式，无需全载内存)")
        print(f"  - 训练集: {len(train_files)} 个文件")
        print(f"  - 验证集: {len(val_files) if val_files else 0} 个文件")
        print(f"  - 测试集: {len(test_files) if test_files else 0} 个文件")

        # 4️⃣ 训练循环
        start_time = time.time()
        skipped_batches = 0
        best_val_f1 = 0.0
        best_checkpoint = None
        
        for epoch in range(num_epochs):
            model.train()
            total_loss = 0.0
            num_batches = 0
            accum_step = 0

            # 用于计算 epoch 级别的指标
            all_preds = []
            all_targets = []
            all_probs = []

            epoch_start_time = time.time()

            progress_bar = tqdm(train_dataloader, desc=f"📚 Epoch {epoch+1}/{num_epochs} ({split_type})")

            for batch_idx, batch in enumerate(progress_bar):
                try:
                    seqs = batch["sequence"]
                    structs = batch["structure"]
                    texts = batch["text"]
                    labels = batch["label"].to(device).unsqueeze(1)  # BCE需[batch,1]

                    with autocast():  # 混合精度
                        # 批量 tokenizer（使用 config 中的减小 max_length）
                        protein_inputs = model.protein_encoder.tokenizer(
                            seqs, return_tensors="pt", padding=True, truncation=True, max_length=512
                        ).to(device)

                        text_inputs = model.text_encoder.tokenizer(
                            texts, return_tensors="pt", padding=True, truncation=True, max_length=256
                        ).to(device)

                        structure_inputs = None
                        if model.structure_config is not None:
                            structure_inputs = model.structure_encoder.tokenizer(
                                structs, return_tensors="pt", padding=True, truncation=True, max_length=512
                            ).to(device)

                        # 前向传播
                        logits = model(
                            protein_inputs=protein_inputs,
                            text_inputs=text_inputs,
                            structure_inputs=structure_inputs
                        )  # 假设输出[batch,1]

                        pred_probs = torch.sigmoid(logits)  # 始终使用sigmoid

                        # 计算损失
                        loss = criterion(logits, labels)

                        l2_reg = sum([p.norm(2).pow(2) for p in model.affinity_head.parameters() if p.requires_grad])  # 只针对 head
                        loss = loss + 1e-3 * l2_reg  # 正则化系数

                    # 梯度累积与 AMP 缩放
                    loss = loss / accumulation_steps
                    scaler.scale(loss).backward()
                    accum_step += 1

                    if accum_step % accumulation_steps == 0:
                        scaler.step(optimizer)
                        scaler.update()
                        optimizer.zero_grad()
                        accum_step = 0

                    batch_loss = loss.item() * accumulation_steps  # 反缩放用于日志

                    total_loss += batch_loss
                    num_batches += 1

                    # 收集预测和真实值（及早 detach 以节省内存）
                    all_preds.extend((pred_probs > 0.5).float().detach().cpu().numpy().flatten())
                    all_probs.extend(pred_probs.detach().cpu().numpy().flatten())
                    all_targets.extend(labels.detach().cpu().numpy().flatten())

                    # 记录 batch 损失（包括当前 lr）
                    current_lr = scheduler.get_last_lr()[0]
                    split_logger.log_batch_loss(epoch + 1, batch_idx, batch_loss, current_lr)

                    # 更新进度条
                    progress_bar.set_postfix({
                        "loss": f"{batch_loss:.4f}", 
                        "lr": f"{current_lr:.6f}",
                        "avg": f"{total_loss/num_batches:.4f}",
                        "skipped_batches": skipped_batches,
                        "mem": f"{torch.cuda.max_memory_allocated() / 1e9:.2f}GB" if 'cuda' in device else "N/A"
                    })

                    # 内存清理
                    del protein_inputs, text_inputs, structure_inputs, logits, pred_probs, loss
                    if 'cuda' in device:
                        torch.cuda.empty_cache()

                except torch.cuda.OutOfMemoryError as e:
                    print(f"⚠️ OOM Batch {batch_idx} 跳过: {str(e)}")
                    skipped_batches += 1
                    optimizer.zero_grad()
                    scaler.update()  # 重置scaler
                    if 'cuda' in device:
                        torch.cuda.empty_cache()
                    continue
                except Exception as e:
                    print(f"⚠️ Batch {batch_idx} 跳过: {str(e)}")
                    traceback.print_exc()
                    skipped_batches += 1
                    optimizer.zero_grad()  # 错误时清零梯度
                    continue

            # 每个 epoch 末尾 scheduler step
            scheduler.step()

            # 清空 CUDA 缓存
            if 'cuda' in device:
                torch.cuda.empty_cache()

            epoch_end_time = time.time()
            epoch_duration = epoch_end_time - epoch_start_time

            if num_batches > 0:
                # 计算 epoch 级别指标
                all_preds = np.array(all_preds)
                all_targets = np.array(all_targets)
                all_probs = np.array(all_probs)

                avg_epoch_loss = total_loss / num_batches

                # 计算指标（二分类）
                if len(all_preds) > 1:
                    acc = accuracy_score(all_targets, all_preds)
                    f1 = f1_score(all_targets, all_preds, average='binary')
                    precision = precision_score(all_targets, all_preds, average='binary')
                    recall = recall_score(all_targets, all_preds, average='binary')
                    auc = roc_auc_score(all_targets, all_probs)
                else:
                    acc, f1, precision, recall, auc = 0, 0, 0, 0, 0 

                # 记录 epoch 统计
                current_lr = scheduler.get_last_lr()[0]
                epoch_metrics = {
                    'loss': avg_epoch_loss,
                    'acc': acc,
                    'f1': f1,
                    'precision': precision,
                    'recall': recall,
                    'auc': auc,
                    'lr': current_lr
                }
                split_logger.log_epoch(epoch + 1, epoch_metrics, current_lr)

                print(f"✅ [Epoch {epoch+1}] 平均Loss = {avg_epoch_loss:.4f}, 当前LR = {current_lr:.6f}, 耗时: {epoch_duration:.1f}s")
                print(f"   📊 详细指标:")
                print(f"   - ACC: {acc:.4f}")
                print(f"   - F1: {f1:.4f}")
                print(f"   - Precision: {precision:.4f}")
                print(f"   - Recall: {recall:.4f}")
                print(f"   - AUC: {auc:.4f}")

                # 验证阶段
                model.eval()
                val_total_loss = 0.0
                val_num_batches = 0
                val_all_preds = []
                val_all_targets = []
                val_all_probs = []

                with torch.no_grad(), autocast():
                    val_progress_bar = tqdm(val_dataloader, desc=f"🔍 Validation {epoch+1}/{num_epochs} ({split_type})")
                    for batch in val_progress_bar:
                        seqs = batch["sequence"]
                        structs = batch["structure"]
                        texts = batch["text"]
                        labels = batch["label"].to(device).unsqueeze(1)

                        protein_inputs = model.protein_encoder.tokenizer(
                            seqs, return_tensors="pt", padding=True, truncation=True, max_length=512
                        ).to(device)

                        text_inputs = model.text_encoder.tokenizer(
                            texts, return_tensors="pt", padding=True, truncation=True, max_length=256
                        ).to(device)

                        structure_inputs = None
                        if model.structure_config is not None:
                            structure_inputs = model.structure_encoder.tokenizer(
                                structs, return_tensors="pt", padding=True, truncation=True, max_length=512
                            ).to(device)

                        logits = model(
                            protein_inputs=protein_inputs,
                            text_inputs=text_inputs,
                            structure_inputs=structure_inputs
                        )

                        pred_probs = torch.sigmoid(logits)  # 始终使用sigmoid

                        loss = criterion(logits, labels)
                        val_total_loss += loss.item()
                        val_num_batches += 1

                        val_all_preds.extend((pred_probs > 0.5).float().detach().cpu().numpy().flatten())
                        val_all_probs.extend(pred_probs.detach().cpu().numpy().flatten())
                        val_all_targets.extend(labels.detach().cpu().numpy().flatten())

                        del protein_inputs, text_inputs, structure_inputs, logits, pred_probs, loss
                        if 'cuda' in device:
                            torch.cuda.empty_cache()

                if 'cuda' in device:
                    torch.cuda.empty_cache()

                if val_num_batches > 0:
                    avg_val_loss = val_total_loss / val_num_batches
                    val_acc = accuracy_score(val_all_targets, val_all_preds)
                    val_f1 = f1_score(val_all_targets, val_all_preds, average='binary')
                    val_precision = precision_score(val_all_targets, val_all_preds, average='binary')
                    val_recall = recall_score(val_all_targets, val_all_preds, average='binary')
                    val_auc = roc_auc_score(val_all_targets, val_all_probs)

                    val_metrics = {
                        'val_loss': avg_val_loss,
                        'val_acc': val_acc,
                        'val_f1': val_f1,
                        'val_precision': val_precision,
                        'val_recall': val_recall,
                        'val_auc': val_auc
                    }
                    split_logger.log_epoch(epoch + 1, val_metrics, current_lr, prefix='val_')

                    print(f"   📊 验证指标:")
                    print(f"   - Val Loss: {avg_val_loss:.4f}")
                    print(f"   - Val ACC: {val_acc:.4f}")
                    print(f"   - Val F1: {val_f1:.4f}")
                    print(f"   - Val Precision: {val_precision:.4f}")
                    print(f"   - Val Recall: {val_recall:.4f}")
                    print(f"   - Val AUC: {val_auc:.4f}")

                    # 保存检查点
                    checkpoint = {
                        'epoch': epoch + 1,
                        'model_state_dict': model.state_dict(),
                        'optimizer_state_dict': optimizer.state_dict(),
                        'scheduler_state_dict': scheduler.state_dict(),
                        'metrics': epoch_metrics,
                        'val_metrics': val_metrics,
                        'config': config
                    }

                    torch.save(checkpoint, os.path.join(split_logger.checkpoint_dir, f"epoch_{epoch+1}.pt"))

                    # 保存最佳模型（基于 val_f1）
                    if val_f1 > best_val_f1:
                        best_val_f1 = val_f1
                        best_checkpoint = checkpoint
                        torch.save(checkpoint, os.path.join(split_logger.checkpoint_dir, "best_model.pt"))
                        print(f"🏆 新的最佳模型已保存，Val F1: {val_f1:.4f}")

            model.train()  # 恢复训练模式（如果循环后有其他操作）

        # 训练结束，进行测试
        if best_checkpoint is not None:
            print(f"🔄 加载最佳模型进行测试 ({split_type})")
            model.load_state_dict(best_checkpoint['model_state_dict'])
        
        model.eval()
        test_total_loss = 0.0
        test_num_batches = 0
        test_all_preds = []
        test_all_targets = []
        test_all_probs = []
        
        with torch.no_grad():
            test_progress_bar = tqdm(test_dataloader, desc=f"🧪 Testing ({split_type})")
            for batch in test_progress_bar:
                seqs = batch["sequence"]
                structs = batch["structure"]
                texts = batch["text"]
                labels = batch["label"].to(device).unsqueeze(1)
                
                protein_inputs = model.protein_encoder.tokenizer(
                    seqs, return_tensors="pt", padding=True, truncation=True, max_length=512
                ).to(device)
                
                text_inputs = model.text_encoder.tokenizer(
                    texts, return_tensors="pt", padding=True, truncation=True, max_length=256
                ).to(device)
                
                structure_inputs = None
                if model.structure_config is not None:
                    structure_inputs = model.structure_encoder.tokenizer(
                        structs, return_tensors="pt", padding=True, truncation=True, max_length=512
                    ).to(device)
                
                logits = model(
                    protein_inputs=protein_inputs,
                    text_inputs=text_inputs,
                    structure_inputs=structure_inputs
                )
                
                pred_probs = torch.sigmoid(logits)  # 始终使用sigmoid
                
                loss = criterion(logits, labels)
                test_total_loss += loss.item()
                test_num_batches += 1
                
                test_all_preds.extend((pred_probs > 0.5).float().detach().cpu().numpy().flatten())
                test_all_probs.extend(pred_probs.detach().cpu().numpy().flatten())
                test_all_targets.extend(labels.detach().cpu().numpy().flatten())
                
                del protein_inputs, text_inputs, structure_inputs, logits, pred_probs, loss
                if 'cuda' in device:
                    torch.cuda.empty_cache()
        
        if test_num_batches > 0:
            avg_test_loss = test_total_loss / test_num_batches
            test_acc = accuracy_score(test_all_targets, test_all_preds)
            test_f1 = f1_score(test_all_targets, test_all_preds, average='binary')
            test_precision = precision_score(test_all_targets, test_all_preds, average='binary')
            test_recall = recall_score(test_all_targets, test_all_preds, average='binary')
            test_auc = roc_auc_score(test_all_targets, test_all_probs)
            
            test_metrics = {
                'test_loss': avg_test_loss,
                'test_acc': test_acc,
                'test_f1': test_f1,
                'test_precision': test_precision,
                'test_recall': test_recall,
                'test_auc': test_auc
            }
            split_logger.log_epoch(num_epochs, test_metrics, 0, prefix='test_')
            
            print(f"   📊 测试指标 ({split_type}):")
            print(f"   - Test Loss: {avg_test_loss:.4f}")
            print(f"   - Test ACC: {test_acc:.4f}")
            print(f"   - Test F1: {test_f1:.4f}")
            print(f"   - Test Precision: {test_precision:.4f}")
            print(f"   - Test Recall: {test_recall:.4f}")
            print(f"   - Test AUC: {test_auc:.4f}")

        # 保存训练总结和图表
        total_training_time = time.time() - start_time
        split_logger.save_training_summary(config, total_training_time)
        split_logger.plot_loss_curves()
        print(f"{PROJECT_NAME}: {split_type} split finished.")
        
        print(f"\n🎉 {split_type} 微调完成!")
        print(f"⏱️ 总训练时间: {total_training_time:.1f} 秒 ({total_training_time/60:.1f} 分钟)")
        print(f"📊 结果保存在: {split_result_dir}")
        print(f"📈 损失图表: {os.path.join(split_logger.plot_dir, 'loss_curves.png')}")
        print(f"📊 指标图表: {os.path.join(split_logger.plot_dir, 'metrics_curves.png')}")
        print(f"💾 最佳模型: {os.path.join(split_logger.checkpoint_dir, 'best_model.pt')}") 
        
if __name__ == "__main__":
    train_function_model()