mirror of
https://github.com/RYDE-WORK/lnp_ml.git
synced 2026-09-18 14:23:20 +08:00
1202 lines
44 KiB
Python
1202 lines
44 KiB
Python
"""
|
||
嵌套交叉验证 + Optuna 超参调优
|
||
|
||
外层 5-fold StratifiedKFold（20% test / 80% train）
|
||
内层 3-fold StratifiedKFold（在 80% 上做 Optuna 超参搜索）
|
||
|
||
使用方法:
|
||
python -m lnp_ml.modeling.nested_cv_optuna
|
||
|
||
或通过 Makefile:
|
||
make nested_cv_tune DEVICE=cuda
|
||
"""
|
||
|
||
import json
|
||
import math
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import Dict, List, Optional, Tuple, Union
|
||
|
||
import numpy as np
|
||
import pandas as pd
|
||
import torch
|
||
from torch.utils.data import DataLoader, Subset
|
||
from sklearn.model_selection import StratifiedKFold
|
||
from loguru import logger
|
||
import typer
|
||
|
||
try:
|
||
import optuna
|
||
from optuna.samplers import TPESampler
|
||
except ImportError:
|
||
optuna = None
|
||
TPESampler = None
|
||
|
||
from lnp_ml.config import MODELS_DIR, INTERIM_DATA_DIR
|
||
from lnp_ml.dataset import (
|
||
LNPDataset,
|
||
collate_fn,
|
||
process_dataframe,
|
||
TARGET_CLASSIFICATION_PDI,
|
||
TARGET_CLASSIFICATION_EE,
|
||
TARGET_TOXIC,
|
||
)
|
||
from tqdm import tqdm
|
||
|
||
from lnp_ml.modeling.models import LNPModel, LNPModelWithoutMPNN
|
||
from lnp_ml.modeling.layers.llm_prompt import DEFAULT_MOLT5_PATH
|
||
from lnp_ml.utils.seed import set_global_seed
|
||
from lnp_ml.modeling.encoders.rdkit_encoder import CachedRDKitEncoder
|
||
from lnp_ml.modeling.trainer_balanced import (
|
||
ClassWeights,
|
||
LossWeightsBalanced,
|
||
compute_class_weights_from_loader,
|
||
train_with_early_stopping,
|
||
train_fixed_epochs,
|
||
validate_balanced,
|
||
)
|
||
|
||
# MPNN ensemble 默认路径
|
||
DEFAULT_MPNN_ENSEMBLE_DIR = MODELS_DIR / "mpnn" / "all_amine_split_for_LiON"
|
||
|
||
app = typer.Typer()
|
||
|
||
|
||
# ============ CompositeStrata 复合分层标签 ============
|
||
|
||
def build_composite_strata(
|
||
df: pd.DataFrame,
|
||
min_stratum_count: int = 5,
|
||
) -> Tuple[np.ndarray, Dict]:
|
||
"""
|
||
构建复合分层标签（toxic × PDI × EE）。
|
||
|
||
Args:
|
||
df: 处理后的 DataFrame
|
||
min_stratum_count: 每个 stratum 最少样本数，低于此值合并为 RARE
|
||
|
||
Returns:
|
||
(strata_array, strata_info)
|
||
- strata_array: 每个样本的 stratum 编码（整数）
|
||
- strata_info: 统计信息
|
||
"""
|
||
n = len(df)
|
||
strata_labels = []
|
||
|
||
for i in range(n):
|
||
# Toxic stratum
|
||
if TARGET_TOXIC in df.columns:
|
||
toxic_val = df[TARGET_TOXIC].iloc[i]
|
||
if pd.notna(toxic_val) and toxic_val >= 0:
|
||
toxic_str = str(int(toxic_val))
|
||
else:
|
||
toxic_str = "NA"
|
||
else:
|
||
toxic_str = "NA"
|
||
|
||
# PDI stratum
|
||
if all(col in df.columns for col in TARGET_CLASSIFICATION_PDI):
|
||
pdi_vals = df[TARGET_CLASSIFICATION_PDI].iloc[i].values
|
||
if pdi_vals.sum() > 0:
|
||
pdi_str = str(int(np.argmax(pdi_vals)))
|
||
else:
|
||
pdi_str = "NA"
|
||
else:
|
||
pdi_str = "NA"
|
||
|
||
# EE stratum
|
||
if all(col in df.columns for col in TARGET_CLASSIFICATION_EE):
|
||
ee_vals = df[TARGET_CLASSIFICATION_EE].iloc[i].values
|
||
if ee_vals.sum() > 0:
|
||
ee_str = str(int(np.argmax(ee_vals)))
|
||
else:
|
||
ee_str = "NA"
|
||
else:
|
||
ee_str = "NA"
|
||
|
||
strata_labels.append(f"T{toxic_str}|P{pdi_str}|E{ee_str}")
|
||
|
||
# 统计各 stratum 的样本数
|
||
unique_strata, counts = np.unique(strata_labels, return_counts=True)
|
||
strata_counts = dict(zip(unique_strata, counts))
|
||
|
||
# 将稀疏 strata 合并为 RARE
|
||
rare_strata = [s for s, c in strata_counts.items() if c < min_stratum_count]
|
||
|
||
final_labels = []
|
||
for label in strata_labels:
|
||
if label in rare_strata:
|
||
final_labels.append("RARE")
|
||
else:
|
||
final_labels.append(label)
|
||
|
||
# 编码为整数
|
||
unique_final, encoded = np.unique(final_labels, return_inverse=True)
|
||
|
||
strata_info = {
|
||
"original_strata_counts": strata_counts,
|
||
"rare_strata": rare_strata,
|
||
"final_strata": list(unique_final),
|
||
"final_strata_counts": dict(zip(*np.unique(final_labels, return_counts=True))),
|
||
"n_rare_merged": sum(strata_counts[s] for s in rare_strata) if rare_strata else 0,
|
||
}
|
||
|
||
logger.info(f"Built composite strata: {len(unique_final)} unique strata")
|
||
logger.info(f" Rare strata merged: {len(rare_strata)} types, {strata_info['n_rare_merged']} samples")
|
||
|
||
return encoded.astype(np.int64), strata_info
|
||
|
||
|
||
# ============ RDKit 缓存预热 ============
|
||
|
||
def warmup_rdkit_cache(smiles_list: List[str], batch_size: int = 256) -> Dict:
|
||
"""预热 RDKit 特征缓存，返回可跨模型共享的缓存字典。"""
|
||
encoder = CachedRDKitEncoder()
|
||
unique_smiles = list(set(smiles_list))
|
||
logger.info(f"Warming up RDKit cache for {len(unique_smiles)} unique SMILES...")
|
||
for i in tqdm(range(0, len(unique_smiles), batch_size), desc="Cache warmup"):
|
||
batch = unique_smiles[i:i + batch_size]
|
||
encoder(batch)
|
||
logger.success(f"Cache warmup complete. Cached {len(encoder._cache)} SMILES.")
|
||
return encoder._cache
|
||
|
||
|
||
# ============ 模型创建 ============
|
||
|
||
def find_mpnn_ensemble_paths(base_dir: Path = DEFAULT_MPNN_ENSEMBLE_DIR) -> List[str]:
|
||
"""自动查找 MPNN ensemble 的 model.pt 文件。"""
|
||
model_paths = sorted(base_dir.glob("cv_*/fold_*/model_*/model.pt"))
|
||
if not model_paths:
|
||
raise FileNotFoundError(f"No model.pt files found in {base_dir}")
|
||
return [str(p) for p in model_paths]
|
||
|
||
|
||
def _load_external_retrieval_pool(csv_path="data/external/all_data_LiON.csv"):
|
||
"""读外部 LiON 数据，筛 mRNA+Mouse，按唯一分子聚合 delivery，返回 (smiles_list, labels[N,1])。"""
|
||
import pandas as pd
|
||
import numpy as np
|
||
df = pd.read_csv(csv_path, low_memory=False)
|
||
sub = df[(df["Cargo_type"] == "mRNA") & (df["Model_type"] == "Mouse")].copy()
|
||
sub["quantified_delivery"] = pd.to_numeric(sub["quantified_delivery"], errors="coerce")
|
||
sub = sub.dropna(subset=["quantified_delivery", "smiles"])
|
||
# 同一分子多器官记录取平均，聚合成唯一分子
|
||
agg = sub.groupby("smiles")["quantified_delivery"].mean()
|
||
smiles_list = agg.index.tolist()
|
||
labels = agg.values.reshape(-1, 1).astype("float32")
|
||
return smiles_list, labels
|
||
|
||
|
||
def create_model(
|
||
d_model: int = 256,
|
||
num_heads: int = 8,
|
||
n_attn_layers: int = 4,
|
||
fusion_strategy: str = "attention",
|
||
head_hidden_dim: int = 128,
|
||
dropout: float = 0.1,
|
||
use_mpnn: bool = False,
|
||
mpnn_device: str = "cpu",
|
||
chemeleon_cache: Optional[str] = None,
|
||
unimol_cache: Optional[str] = None,
|
||
moleculestm_cache: Optional[str] = None,
|
||
mole_cache: Optional[str] = None,
|
||
set_transformer_block: str = "sab",
|
||
# MoE
|
||
use_moe: bool = False,
|
||
moe_n_experts: int = 4,
|
||
moe_top_k: int = 2,
|
||
moe_expert_hidden_mult: int = 2,
|
||
moe_jitter_noise: float = 0.0,
|
||
# LLM（打包成 dict 以减少多进程参数传递）
|
||
llm_kwargs: Optional[Dict] = None,
|
||
# 检索增强
|
||
use_retrieval: bool = False,
|
||
retr_feature_dim: int = 0,
|
||
) -> Union[LNPModel, LNPModelWithoutMPNN]:
|
||
"""创建模型。llm_kwargs 含 use_llm / llm_model_path / llm_freeze / llm_use_lora / llm_lora_*。"""
|
||
extra_kwargs = dict(
|
||
set_transformer_block=set_transformer_block,
|
||
use_moe=use_moe,
|
||
moe_n_experts=moe_n_experts,
|
||
moe_top_k=moe_top_k,
|
||
moe_expert_hidden_mult=moe_expert_hidden_mult,
|
||
moe_jitter_noise=moe_jitter_noise,
|
||
use_retrieval=use_retrieval,
|
||
retr_feature_dim=retr_feature_dim,
|
||
chemeleon_cache_path=chemeleon_cache,
|
||
unimol_cache_path=unimol_cache,
|
||
moleculestm_cache_path=moleculestm_cache,
|
||
mole_cache_path=mole_cache,
|
||
**(llm_kwargs or {}),
|
||
)
|
||
|
||
if use_mpnn:
|
||
ensemble_paths = find_mpnn_ensemble_paths()
|
||
return LNPModel(
|
||
d_model=d_model,
|
||
num_heads=num_heads,
|
||
n_attn_layers=n_attn_layers,
|
||
fusion_strategy=fusion_strategy,
|
||
head_hidden_dim=head_hidden_dim,
|
||
dropout=dropout,
|
||
mpnn_ensemble_paths=ensemble_paths,
|
||
mpnn_device=mpnn_device,
|
||
**extra_kwargs,
|
||
)
|
||
else:
|
||
return LNPModelWithoutMPNN(
|
||
d_model=d_model,
|
||
num_heads=num_heads,
|
||
n_attn_layers=n_attn_layers,
|
||
fusion_strategy=fusion_strategy,
|
||
head_hidden_dim=head_hidden_dim,
|
||
dropout=dropout,
|
||
**extra_kwargs,
|
||
)
|
||
|
||
|
||
def _build_rag_pool(full_dataset, idx):
|
||
"""从 dataset + 索引构造 RAG 池：返回 (smiles, delivery[N], extra_labels)。"""
|
||
import numpy as np
|
||
idx = np.asarray(idx)
|
||
smiles = [full_dataset.smiles[i] for i in idx]
|
||
delivery = full_dataset.delivery[idx].reshape(-1)
|
||
extra = {}
|
||
if full_dataset.size is not None:
|
||
extra["size"] = (full_dataset.size[idx], ~np.isnan(full_dataset.size[idx]))
|
||
if full_dataset.pdi is not None:
|
||
extra["pdi"] = (full_dataset.pdi[idx], full_dataset.pdi_valid[idx])
|
||
if full_dataset.ee is not None:
|
||
extra["ee"] = (full_dataset.ee[idx], full_dataset.ee_valid[idx])
|
||
if full_dataset.toxic is not None:
|
||
extra["toxic"] = (full_dataset.toxic[idx], full_dataset.toxic[idx] >= 0)
|
||
if full_dataset.biodist is not None:
|
||
extra["biodist"] = (full_dataset.biodist[idx], full_dataset.biodist_valid[idx])
|
||
return smiles, delivery, extra
|
||
|
||
|
||
# ============ 评估指标 ============
|
||
|
||
def evaluate_on_test(
|
||
model: torch.nn.Module,
|
||
test_loader: DataLoader,
|
||
device: torch.device,
|
||
) -> Dict:
|
||
"""在测试集上评估模型"""
|
||
from scipy.special import rel_entr
|
||
from sklearn.metrics import (
|
||
mean_squared_error,
|
||
mean_absolute_error,
|
||
r2_score,
|
||
accuracy_score,
|
||
precision_score,
|
||
recall_score,
|
||
f1_score,
|
||
)
|
||
|
||
model.eval()
|
||
|
||
preds = {
|
||
"size": [], "delivery": [], "pdi": [], "ee": [], "toxic": [], "biodist": []
|
||
}
|
||
targets = {
|
||
"size": [], "delivery": [], "pdi": [], "ee": [], "toxic": [], "biodist": []
|
||
}
|
||
|
||
with torch.no_grad():
|
||
for batch in test_loader:
|
||
smiles = batch["smiles"]
|
||
tabular = {k: v.to(device) for k, v in batch["tabular"].items()}
|
||
tgts = batch["targets"]
|
||
masks = batch["mask"]
|
||
|
||
outputs = model(smiles, tabular)
|
||
|
||
# 收集预测和真实值
|
||
for task in ["size", "delivery"]:
|
||
if task in masks and masks[task].any():
|
||
m = masks[task]
|
||
key = task if task == "size" else "delivery"
|
||
preds[task].extend(outputs[key].squeeze(-1)[m].cpu().numpy().tolist())
|
||
targets[task].extend(tgts[key][m].cpu().numpy().tolist())
|
||
|
||
for task in ["pdi", "ee", "toxic"]:
|
||
if task in masks and masks[task].any():
|
||
m = masks[task]
|
||
preds[task].extend(outputs[task][m].argmax(dim=-1).cpu().numpy().tolist())
|
||
targets[task].extend(tgts[task][m].cpu().numpy().tolist())
|
||
|
||
if "biodist" in masks and masks["biodist"].any():
|
||
m = masks["biodist"]
|
||
preds["biodist"].extend(outputs["biodist"][m].cpu().numpy().tolist())
|
||
targets["biodist"].extend(tgts["biodist"][m].cpu().numpy().tolist())
|
||
|
||
# 计算指标
|
||
results = {}
|
||
|
||
# 回归任务
|
||
for task in ["size", "delivery"]:
|
||
if preds[task]:
|
||
p = np.array(preds[task])
|
||
t = np.array(targets[task])
|
||
results[task] = {
|
||
"n_samples": len(p),
|
||
"mse": float(mean_squared_error(t, p)),
|
||
"rmse": float(np.sqrt(mean_squared_error(t, p))),
|
||
"mae": float(mean_absolute_error(t, p)),
|
||
"r2": float(r2_score(t, p)),
|
||
}
|
||
|
||
# 分类任务
|
||
for task in ["pdi", "ee", "toxic"]:
|
||
if preds[task]:
|
||
p = np.array(preds[task])
|
||
t = np.array(targets[task])
|
||
results[task] = {
|
||
"n_samples": len(p),
|
||
"accuracy": float(accuracy_score(t, p)),
|
||
"precision": float(precision_score(t, p, average="macro", zero_division=0)),
|
||
"recall": float(recall_score(t, p, average="macro", zero_division=0)),
|
||
"f1": float(f1_score(t, p, average="macro", zero_division=0)),
|
||
}
|
||
|
||
# 分布任务
|
||
if preds["biodist"]:
|
||
p = np.array(preds["biodist"])
|
||
t = np.array(targets["biodist"])
|
||
|
||
def kl_divergence(p_arr, q_arr, eps=1e-10):
|
||
p_arr = np.clip(p_arr, eps, 1.0)
|
||
q_arr = np.clip(q_arr, eps, 1.0)
|
||
return float(np.sum(rel_entr(p_arr, q_arr), axis=-1).mean())
|
||
|
||
def js_divergence(p_arr, q_arr, eps=1e-10):
|
||
p_arr = np.clip(p_arr, eps, 1.0)
|
||
q_arr = np.clip(q_arr, eps, 1.0)
|
||
m = 0.5 * (p_arr + q_arr)
|
||
return float(0.5 * (np.sum(rel_entr(p_arr, m), axis=-1) + np.sum(rel_entr(q_arr, m), axis=-1)).mean())
|
||
|
||
results["biodist"] = {
|
||
"n_samples": len(p),
|
||
"kl_divergence": kl_divergence(t, p),
|
||
"js_divergence": js_divergence(t, p),
|
||
}
|
||
|
||
return results
|
||
|
||
|
||
# ============ 预训练权重加载 ============
|
||
|
||
def load_pretrain_weights_to_model(
|
||
model: Union[LNPModel, LNPModelWithoutMPNN],
|
||
pretrain_state_dict: Dict,
|
||
d_model: int,
|
||
pretrain_config: Dict,
|
||
load_delivery_head: bool = True,
|
||
) -> bool:
|
||
"""
|
||
加载预训练权重到模型。
|
||
|
||
Returns:
|
||
是否成功加载
|
||
"""
|
||
if pretrain_config.get("d_model") != d_model:
|
||
logger.warning(
|
||
f"d_model mismatch: pretrain={pretrain_config.get('d_model')}, "
|
||
f"current={d_model}. Skipping pretrain loading."
|
||
)
|
||
return False
|
||
|
||
model.load_pretrain_weights(
|
||
pretrain_state_dict=pretrain_state_dict,
|
||
load_delivery_head=load_delivery_head,
|
||
strict=False,
|
||
)
|
||
return True
|
||
|
||
|
||
# ============ 内层 Optuna 调参 ============
|
||
def run_inner_optuna(
|
||
full_dataset: LNPDataset,
|
||
inner_train_indices: np.ndarray,
|
||
strata: np.ndarray,
|
||
device: torch.device,
|
||
n_trials: int = 20,
|
||
epochs_per_trial: int = 30,
|
||
patience: int = 10,
|
||
batch_size: int = 32,
|
||
n_inner_folds: int = 3,
|
||
use_mpnn: bool = False,
|
||
chemeleon_cache: Optional[str] = None,
|
||
unimol_cache: Optional[str] = None,
|
||
moleculestm_cache: Optional[str] = None,
|
||
mole_cache: Optional[str] = None,
|
||
seed: int = 42,
|
||
study_path: Optional[Path] = None,
|
||
pretrain_state_dict: Optional[Dict] = None,
|
||
pretrain_config: Optional[Dict] = None,
|
||
load_delivery_head: bool = True,
|
||
rdkit_cache: Optional[Dict] = None,
|
||
use_moe: bool = False,
|
||
moe_n_experts: int = 4,
|
||
moe_top_k: int = 2,
|
||
moe_expert_hidden_mult: int = 2,
|
||
moe_jitter_noise: float = 0.0,
|
||
set_transformer_block: str = "sab",
|
||
llm_kwargs: Optional[Dict] = None,
|
||
) -> Tuple[Dict, int, optuna.Study]:
|
||
"""
|
||
在内层数据上运行 Optuna 超参搜索。
|
||
|
||
Args:
|
||
full_dataset: 完整数据集
|
||
inner_train_indices: 内层训练数据的索引（相对于 full_dataset）
|
||
strata: 每个样本的分层标签
|
||
device: 设备
|
||
n_trials: Optuna 试验数
|
||
epochs_per_trial: 每个试验的最大 epoch
|
||
patience: 早停耐心值
|
||
batch_size: 批次大小
|
||
n_inner_folds: 内层折数
|
||
use_mpnn: 是否使用 MPNN
|
||
seed: 随机种子
|
||
study_path: 可选的 study 持久化路径
|
||
pretrain_state_dict: 预训练权重
|
||
pretrain_config: 预训练配置
|
||
load_delivery_head: 是否加载 delivery head 权重
|
||
|
||
Returns:
|
||
(best_params, epoch_mean, study)
|
||
"""
|
||
if optuna is None:
|
||
raise ImportError("Optuna not installed. Run: pip install optuna")
|
||
|
||
inner_strata = strata[inner_train_indices]
|
||
|
||
# 固定架构参数（与预训练一致，确保权重完整加载）
|
||
_cfg = pretrain_config or {}
|
||
fixed_d_model = _cfg.get("d_model", 256)
|
||
fixed_num_heads = _cfg.get("num_heads", 8)
|
||
fixed_n_attn_layers = _cfg.get("n_attn_layers", 4)
|
||
fixed_fusion_strategy = _cfg.get("fusion_strategy", "attention")
|
||
fixed_head_hidden_dim = _cfg.get("head_hidden_dim", 128)
|
||
logger.info(
|
||
f"Fixed architecture params: d_model={fixed_d_model}, num_heads={fixed_num_heads}, "
|
||
f"n_attn_layers={fixed_n_attn_layers}, fusion={fixed_fusion_strategy}, "
|
||
f"head_hidden_dim={fixed_head_hidden_dim}"
|
||
)
|
||
|
||
def objective(trial: optuna.Trial) -> float:
|
||
d_model = fixed_d_model
|
||
num_heads = fixed_num_heads
|
||
n_attn_layers = fixed_n_attn_layers
|
||
fusion_strategy = fixed_fusion_strategy
|
||
head_hidden_dim = fixed_head_hidden_dim
|
||
|
||
# 搜索训练超参数
|
||
dropout = trial.suggest_float("dropout", 0.1, 0.5)
|
||
lr = trial.suggest_float("lr", 1e-5, 1e-3, log=True)
|
||
weight_decay = trial.suggest_float("weight_decay", 1e-5, 1e-1, log=True)
|
||
backbone_lr_ratio = trial.suggest_float("backbone_lr_ratio", 0.01, 1.0, log=True)
|
||
|
||
# MoE / LLM
|
||
use_llm = bool(llm_kwargs.get("use_llm", False))
|
||
moe_ne_t = trial.suggest_categorical("moe_n_experts", [2, 4]) if use_moe else moe_n_experts
|
||
moe_tk_t = trial.suggest_int("moe_top_k", 1, 2) if use_moe else moe_top_k
|
||
moe_hm_t = trial.suggest_categorical("moe_expert_hidden_mult", [1, 2]) if use_moe else moe_expert_hidden_mult
|
||
llm_kwargs_t = dict(llm_kwargs)
|
||
if use_llm:
|
||
llm_kwargs_t["llm_use_lora"] = True
|
||
llm_kwargs_t["llm_lora_r"] = trial.suggest_categorical("llm_lora_r", [8, 16, 32])
|
||
|
||
# 内层 3-fold CV
|
||
inner_cv = StratifiedKFold(
|
||
n_splits=n_inner_folds, shuffle=True, random_state=seed
|
||
)
|
||
|
||
fold_val_losses = []
|
||
fold_best_epochs = []
|
||
|
||
for inner_fold, (inner_train_idx, inner_val_idx) in enumerate(
|
||
inner_cv.split(inner_train_indices, inner_strata)
|
||
):
|
||
# 获取实际的数据集索引
|
||
actual_train_idx = inner_train_indices[inner_train_idx]
|
||
actual_val_idx = inner_train_indices[inner_val_idx]
|
||
|
||
# 创建 DataLoader
|
||
train_subset = Subset(full_dataset, actual_train_idx.tolist())
|
||
val_subset = Subset(full_dataset, actual_val_idx.tolist())
|
||
|
||
train_loader = DataLoader(
|
||
train_subset, batch_size=batch_size, shuffle=True,
|
||
collate_fn=collate_fn, drop_last=True,
|
||
)
|
||
val_loader = DataLoader(
|
||
val_subset, batch_size=batch_size, shuffle=False, collate_fn=collate_fn
|
||
)
|
||
|
||
# 计算类权重
|
||
class_weights = compute_class_weights_from_loader(train_loader)
|
||
|
||
# 创建模型
|
||
model = create_model(
|
||
d_model=d_model,
|
||
num_heads=num_heads,
|
||
n_attn_layers=n_attn_layers,
|
||
fusion_strategy=fusion_strategy,
|
||
head_hidden_dim=head_hidden_dim,
|
||
dropout=dropout,
|
||
use_mpnn=use_mpnn,
|
||
mpnn_device=device.type,
|
||
chemeleon_cache=chemeleon_cache,
|
||
unimol_cache=unimol_cache,
|
||
moleculestm_cache=moleculestm_cache,
|
||
mole_cache=mole_cache,
|
||
use_moe=use_moe,
|
||
moe_n_experts=moe_ne_t,
|
||
moe_top_k=moe_tk_t,
|
||
moe_expert_hidden_mult=moe_hm_t,
|
||
moe_jitter_noise=moe_jitter_noise,
|
||
set_transformer_block=set_transformer_block,
|
||
llm_kwargs=llm_kwargs_t,
|
||
)
|
||
if rdkit_cache is not None:
|
||
model.rdkit_encoder._cache = rdkit_cache
|
||
|
||
# 内层 RAG 池：只用内层训练集（防 inner-val 泄漏）
|
||
if llm_kwargs_t.get("use_rag", False) and getattr(model, "llm_prompt", None) is not None:
|
||
_s, _d, _ex = _build_rag_pool(full_dataset, actual_train_idx)
|
||
model.llm_prompt.set_retrieval_pool(
|
||
_s, _d, pool_id=f"inner{inner_fold}", extra_labels=_ex)
|
||
|
||
# 加载预训练权重
|
||
if pretrain_state_dict is not None and pretrain_config is not None:
|
||
load_pretrain_weights_to_model(
|
||
model, pretrain_state_dict, d_model, pretrain_config, load_delivery_head
|
||
)
|
||
|
||
# 训练（带早停）
|
||
result = train_with_early_stopping(
|
||
model=model,
|
||
train_loader=train_loader,
|
||
val_loader=val_loader,
|
||
device=device,
|
||
lr=lr,
|
||
weight_decay=weight_decay,
|
||
epochs=epochs_per_trial,
|
||
patience=patience,
|
||
class_weights=class_weights,
|
||
backbone_lr_ratio=backbone_lr_ratio,
|
||
)
|
||
|
||
fold_val_losses.append(result["best_val_loss"])
|
||
fold_best_epochs.append(result["best_epoch"])
|
||
|
||
# 记录 epoch_mean 到 trial
|
||
epoch_mean = int(round(np.mean(fold_best_epochs)))
|
||
trial.set_user_attr("epoch_mean", epoch_mean)
|
||
trial.set_user_attr("fold_best_epochs", fold_best_epochs)
|
||
|
||
return np.mean(fold_val_losses)
|
||
|
||
# 创建 study
|
||
storage = None
|
||
if study_path is not None:
|
||
storage = f"sqlite:///{study_path}"
|
||
|
||
study = optuna.create_study(
|
||
direction="minimize",
|
||
sampler=TPESampler(seed=seed),
|
||
storage=storage,
|
||
study_name="inner_optuna",
|
||
load_if_exists=True,
|
||
)
|
||
|
||
study.optimize(objective, n_trials=n_trials, show_progress_bar=True)
|
||
|
||
best_params = dict(study.best_trial.params)
|
||
best_params.update({
|
||
"d_model": fixed_d_model,
|
||
"num_heads": fixed_num_heads,
|
||
"n_attn_layers": fixed_n_attn_layers,
|
||
"fusion_strategy": fixed_fusion_strategy,
|
||
"head_hidden_dim": fixed_head_hidden_dim,
|
||
"set_transformer_block": set_transformer_block,
|
||
})
|
||
epoch_mean = study.best_trial.user_attrs.get("epoch_mean", epochs_per_trial)
|
||
|
||
logger.info(f"Best trial: {study.best_trial.number}")
|
||
logger.info(f"Best val_loss: {study.best_trial.value:.4f}")
|
||
logger.info(f"Best params: {best_params}")
|
||
logger.info(f"Epoch mean: {epoch_mean}")
|
||
|
||
return best_params, epoch_mean, study
|
||
|
||
|
||
# ============ 单 fold 执行（可跨进程调用） ============
|
||
def _run_single_outer_fold(
|
||
outer_fold: int,
|
||
outer_train_idx: np.ndarray,
|
||
outer_test_idx: np.ndarray,
|
||
df: pd.DataFrame,
|
||
strata: np.ndarray,
|
||
fold_dir: Path,
|
||
n_trials: int,
|
||
epochs_per_trial: int,
|
||
inner_patience: int,
|
||
batch_size: int,
|
||
n_inner_folds: int,
|
||
use_mpnn: bool,
|
||
chemeleon_cache: Optional[str],
|
||
unimol_cache: Optional[str],
|
||
moleculestm_cache: Optional[str],
|
||
mole_cache: Optional[str],
|
||
seed: int,
|
||
pretrain_state_dict: Optional[Dict],
|
||
pretrain_config: Optional[Dict],
|
||
load_delivery_head: bool,
|
||
device_str: str,
|
||
use_retrieval: bool = False,
|
||
retrieval_source: str = "internal",
|
||
use_moe: bool = False,
|
||
moe_n_experts: int = 4,
|
||
moe_top_k: int = 2,
|
||
moe_expert_hidden_mult: int = 2,
|
||
moe_jitter_noise: float = 0.0,
|
||
set_transformer_block: str = "sab",
|
||
llm_kwargs: Optional[Dict] = None,
|
||
precomputed_best_params: Optional[Dict] = None,
|
||
precomputed_epoch_mean: Optional[int] = None,
|
||
) -> Dict:
|
||
"""
|
||
执行单个外层 fold 的完整流程（内层调参 + 外层训练 + 评估）。
|
||
|
||
所有参数均为可序列化类型，以支持 spawn 多进程。
|
||
"""
|
||
device = torch.device(device_str)
|
||
set_global_seed(seed + outer_fold)
|
||
fold_dir = Path(fold_dir)
|
||
fold_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
# ===== 断点续跑 =====
|
||
_f_metrics = fold_dir / "test_metrics.json"
|
||
_f_bp = fold_dir / "best_params.json"
|
||
_f_em = fold_dir / "epoch_mean.json"
|
||
# (1) 整折已完成 → 读回结果直接跳过
|
||
if _f_metrics.exists() and _f_bp.exists() and _f_em.exists():
|
||
logger.info(f"[RESUME] Outer fold {outer_fold} 已完成，跳过。")
|
||
with open(_f_metrics) as f:
|
||
_tm = json.load(f)
|
||
with open(_f_bp) as f:
|
||
_bp = json.load(f)
|
||
with open(_f_em) as f:
|
||
_em = json.load(f)["epoch_mean"]
|
||
return {"fold": outer_fold, "best_params": _bp,
|
||
"epoch_mean": _em, "test_metrics": _tm}
|
||
# (2) 内层已完成、外层中断 → 复用超参，跳过内层 Optuna
|
||
if precomputed_best_params is None and _f_bp.exists() and _f_em.exists():
|
||
with open(_f_bp) as f:
|
||
precomputed_best_params = json.load(f)
|
||
with open(_f_em) as f:
|
||
precomputed_epoch_mean = json.load(f)["epoch_mean"]
|
||
logger.info(f"[RESUME] fold {outer_fold} 复用已存 best_params，跳过内层 Optuna。")
|
||
|
||
full_dataset = LNPDataset(df)
|
||
# === 断点续跑:已完成的 fold 直接跳过(读回磁盘结果)===
|
||
_tm = fold_dir / "test_metrics.json"
|
||
_bp = fold_dir / "best_params.json"
|
||
_em = fold_dir / "epoch_mean.json"
|
||
if precomputed_best_params is None and _tm.exists() and _bp.exists() and _em.exists():
|
||
logger.success(f"[SKIP] Outer fold {outer_fold} already done, loading cached results.")
|
||
with open(_tm) as _f: _tmd = json.load(_f)
|
||
with open(_bp) as _f: _bpd = json.load(_f)
|
||
with open(_em) as _f: _emd = json.load(_f)
|
||
return {
|
||
"fold": outer_fold,
|
||
"best_params": _bpd,
|
||
"epoch_mean": int(_emd.get("epoch_mean", _emd) if isinstance(_emd, dict) else _emd),
|
||
"test_metrics": _tmd,
|
||
}
|
||
|
||
logger.info(f"\n{'='*60}")
|
||
logger.info(f"OUTER FOLD {outer_fold}")
|
||
logger.info(f"{'='*60}")
|
||
logger.info(f"Train: {len(outer_train_idx)}, Test: {len(outer_test_idx)}")
|
||
|
||
# 预热 RDKit 缓存（在整个 fold 内共享）
|
||
rdkit_cache = warmup_rdkit_cache(full_dataset.smiles)
|
||
|
||
# 保存 split indices
|
||
with open(fold_dir / "splits.json", "w") as f:
|
||
json.dump({
|
||
"outer_train_idx": outer_train_idx.tolist(),
|
||
"outer_test_idx": outer_test_idx.tolist(),
|
||
}, f)
|
||
|
||
# 内层 Optuna 调参（方式 B：已有超参则跳过，直接复用）
|
||
if precomputed_best_params is not None and precomputed_epoch_mean is not None:
|
||
logger.info("Reusing precomputed best_params (skip inner Optuna).")
|
||
best_params = dict(precomputed_best_params)
|
||
epoch_mean = int(precomputed_epoch_mean)
|
||
else:
|
||
logger.info(f"\nRunning inner Optuna with {n_trials} trials...")
|
||
study_path = fold_dir / "optuna_study.sqlite3"
|
||
best_params, epoch_mean, study = run_inner_optuna(
|
||
full_dataset=full_dataset,
|
||
inner_train_indices=outer_train_idx,
|
||
strata=strata,
|
||
device=device,
|
||
n_trials=n_trials,
|
||
epochs_per_trial=epochs_per_trial,
|
||
patience=inner_patience,
|
||
batch_size=batch_size,
|
||
n_inner_folds=n_inner_folds,
|
||
use_mpnn=use_mpnn,
|
||
chemeleon_cache=chemeleon_cache,
|
||
unimol_cache=unimol_cache,
|
||
moleculestm_cache=moleculestm_cache,
|
||
mole_cache=mole_cache,
|
||
seed=seed + outer_fold,
|
||
study_path=study_path,
|
||
pretrain_state_dict=pretrain_state_dict,
|
||
pretrain_config=pretrain_config,
|
||
load_delivery_head=load_delivery_head,
|
||
rdkit_cache=rdkit_cache,
|
||
use_moe=use_moe,
|
||
moe_n_experts=moe_n_experts,
|
||
moe_top_k=moe_top_k,
|
||
moe_expert_hidden_mult=moe_expert_hidden_mult,
|
||
moe_jitter_noise=moe_jitter_noise,
|
||
set_transformer_block=set_transformer_block,
|
||
llm_kwargs=llm_kwargs,
|
||
)
|
||
|
||
# 保存最佳参数
|
||
with open(fold_dir / "best_params.json", "w") as f:
|
||
json.dump(best_params, f, indent=2)
|
||
with open(fold_dir / "epoch_mean.json", "w") as f:
|
||
json.dump({"epoch_mean": epoch_mean}, f)
|
||
|
||
# 外层训练（使用最优超参，固定 epoch 数，不 early-stop）
|
||
logger.info(f"\nTraining outer fold with best params, epochs={epoch_mean}...")
|
||
|
||
# 从 outer_train 再切 10% 作为"监控用 val"（只画曲线、看 gap，不参与选模型/早停）
|
||
_rng = np.random.RandomState(seed + outer_fold)
|
||
_perm = _rng.permutation(len(outer_train_idx))
|
||
_n_val = max(1, int(0.1 * len(outer_train_idx)))
|
||
fit_idx = outer_train_idx[_perm[_n_val:]]
|
||
monitor_idx = outer_train_idx[_perm[:_n_val]]
|
||
|
||
train_subset = Subset(full_dataset, fit_idx.tolist())
|
||
monitor_subset = Subset(full_dataset, monitor_idx.tolist())
|
||
test_subset = Subset(full_dataset, outer_test_idx.tolist())
|
||
|
||
train_loader = DataLoader(
|
||
train_subset, batch_size=batch_size, shuffle=True,
|
||
collate_fn=collate_fn, drop_last=True,
|
||
)
|
||
monitor_loader = DataLoader(
|
||
monitor_subset, batch_size=batch_size, shuffle=False, collate_fn=collate_fn
|
||
)
|
||
test_loader = DataLoader(
|
||
test_subset, batch_size=batch_size, shuffle=False, collate_fn=collate_fn
|
||
)
|
||
|
||
class_weights = compute_class_weights_from_loader(train_loader)
|
||
|
||
model = create_model(
|
||
d_model=best_params["d_model"],
|
||
num_heads=best_params["num_heads"],
|
||
n_attn_layers=best_params["n_attn_layers"],
|
||
fusion_strategy=best_params["fusion_strategy"],
|
||
head_hidden_dim=best_params["head_hidden_dim"],
|
||
dropout=best_params["dropout"],
|
||
use_mpnn=use_mpnn,
|
||
mpnn_device=device.type,
|
||
chemeleon_cache=chemeleon_cache,
|
||
unimol_cache=unimol_cache,
|
||
moleculestm_cache=moleculestm_cache,
|
||
mole_cache=mole_cache,
|
||
use_moe=use_moe,
|
||
moe_n_experts=best_params.get("moe_n_experts", moe_n_experts),
|
||
moe_top_k=best_params.get("moe_top_k", moe_top_k),
|
||
moe_expert_hidden_mult=best_params.get("moe_expert_hidden_mult", moe_expert_hidden_mult),
|
||
moe_jitter_noise=moe_jitter_noise,
|
||
set_transformer_block=best_params.get("set_transformer_block", "sab"),
|
||
llm_kwargs={**llm_kwargs, **({"llm_use_lora": True, "llm_lora_r": best_params["llm_lora_r"]}
|
||
if (llm_kwargs.get("use_llm", False) and "llm_lora_r" in best_params) else {})},
|
||
use_retrieval=use_retrieval,
|
||
retr_feature_dim=(3 if use_retrieval else 0),
|
||
)
|
||
model.rdkit_encoder._cache = rdkit_cache
|
||
if use_retrieval:
|
||
from lnp_ml.modeling.retrieval import MorganRetriever
|
||
if retrieval_source == "external":
|
||
_tr_smiles, _tr_labels = _load_external_retrieval_pool()
|
||
logger.info(f"[retrieval] 外部检索池(mRNA+Mouse)={len(_tr_smiles)} 分子")
|
||
else:
|
||
_tr_smiles = [full_dataset.smiles[i] for i in outer_train_idx]
|
||
_tr_labels = full_dataset.delivery[outer_train_idx].reshape(-1, 1)
|
||
logger.info(f"[retrieval] 内部检索池={len(_tr_smiles)} 训练分子（不含测试集）")
|
||
model.retriever = MorganRetriever(_tr_smiles, _tr_labels, k=5)
|
||
|
||
# ===== RAG 检索池设置（防泄漏：只用训练集分子+标签）=====
|
||
if llm_kwargs and llm_kwargs.get("use_rag", False) and getattr(model, "llm_prompt", None) is not None:
|
||
_s, _d, _ex = _build_rag_pool(full_dataset, outer_train_idx)
|
||
model.llm_prompt.set_retrieval_pool(
|
||
_s, _d, pool_id=f"fold{outer_fold}", extra_labels=_ex)
|
||
logger.info(f"[RAG] 检索池={len(_s)} 训练分子（多任务标签，不含测试集，防泄漏）")
|
||
|
||
if pretrain_state_dict is not None and pretrain_config is not None:
|
||
loaded = load_pretrain_weights_to_model(
|
||
model, pretrain_state_dict, best_params["d_model"],
|
||
pretrain_config, load_delivery_head
|
||
)
|
||
if loaded:
|
||
logger.info(f"Loaded pretrain weights for outer fold {outer_fold}")
|
||
|
||
train_result = train_fixed_epochs(
|
||
model=model,
|
||
train_loader=train_loader,
|
||
val_loader=monitor_loader,
|
||
device=device,
|
||
lr=best_params["lr"],
|
||
weight_decay=best_params["weight_decay"],
|
||
epochs=epoch_mean,
|
||
class_weights=class_weights,
|
||
use_cosine_annealing=True,
|
||
backbone_lr_ratio=best_params.get("backbone_lr_ratio", 1.0),
|
||
freeze_backbone_epochs=3,
|
||
)
|
||
|
||
model.load_state_dict(train_result["final_state"], strict=False)
|
||
model = model.to(device)
|
||
|
||
config = {
|
||
"d_model": best_params["d_model"],
|
||
"num_heads": best_params["num_heads"],
|
||
"n_attn_layers": best_params["n_attn_layers"],
|
||
"fusion_strategy": best_params["fusion_strategy"],
|
||
"head_hidden_dim": best_params["head_hidden_dim"],
|
||
"set_transformer_block": best_params.get("set_transformer_block", "sab"),
|
||
"dropout": best_params["dropout"],
|
||
"use_mpnn": use_mpnn,
|
||
"use_chemeleon": chemeleon_cache is not None,
|
||
"chemeleon_cache": chemeleon_cache,
|
||
"use_unimol": unimol_cache is not None,
|
||
"unimol_cache": unimol_cache,
|
||
"use_moleculestm": moleculestm_cache is not None,
|
||
"moleculestm_cache": moleculestm_cache,
|
||
"use_mole": mole_cache is not None,
|
||
"mole_cache": mole_cache,
|
||
"use_moe": use_moe,
|
||
"moe_n_experts": moe_n_experts,
|
||
"moe_top_k": moe_top_k,
|
||
"moe_expert_hidden_mult": moe_expert_hidden_mult,
|
||
"moe_jitter_noise": moe_jitter_noise,
|
||
**(llm_kwargs or {}),
|
||
}
|
||
|
||
_full_state = train_result["final_state"]
|
||
_slim_state = {k: v for k, v in _full_state.items()
|
||
if not k.startswith("llm_prompt.encoder")}
|
||
torch.save({
|
||
"model_state_dict": _slim_state,
|
||
"config": config,
|
||
"epoch_mean": epoch_mean,
|
||
"best_params": best_params,
|
||
}, fold_dir / "model.pt")
|
||
|
||
with open(fold_dir / "history.json", "w") as f:
|
||
json.dump(train_result["history"], f, indent=2)
|
||
|
||
# 在测试集上评估
|
||
logger.info("Evaluating on outer test set...")
|
||
test_metrics = evaluate_on_test(model, test_loader, device)
|
||
|
||
with open(fold_dir / "test_metrics.json", "w") as f:
|
||
json.dump(test_metrics, f, indent=2)
|
||
|
||
logger.info(f"\nOuter Fold {outer_fold} Test Results:")
|
||
for task, metrics in test_metrics.items():
|
||
if "rmse" in metrics:
|
||
logger.info(f" {task}: RMSE={metrics['rmse']:.4f}, R²={metrics['r2']:.4f}")
|
||
elif "accuracy" in metrics:
|
||
logger.info(f" {task}: Acc={metrics['accuracy']:.4f}, F1={metrics['f1']:.4f}")
|
||
elif "kl_divergence" in metrics:
|
||
logger.info(f" {task}: KL={metrics['kl_divergence']:.4f}, JS={metrics['js_divergence']:.4f}")
|
||
|
||
return {
|
||
"fold": outer_fold,
|
||
"best_params": best_params,
|
||
"epoch_mean": epoch_mean,
|
||
"test_metrics": test_metrics,
|
||
}
|
||
|
||
|
||
# ============ 主流程 ============
|
||
@app.command()
|
||
def main(
|
||
input_path: Path = INTERIM_DATA_DIR / "internal.csv",
|
||
output_dir: Path = MODELS_DIR / "nested_cv",
|
||
resume_dir: Optional[Path] = None,
|
||
# CV 参数
|
||
n_outer_folds: int = 5,
|
||
n_inner_folds: int = 3,
|
||
min_stratum_count: int = 5,
|
||
seed: int = 42,
|
||
# Optuna 参数
|
||
n_trials: int = 20,
|
||
epochs_per_trial: int = 30,
|
||
inner_patience: int = 10,
|
||
# 训练参数
|
||
batch_size: int = 32,
|
||
# 预训练权重
|
||
init_from_pretrain: Optional[Path] = None,
|
||
load_delivery_head: bool = False,
|
||
# MPNN
|
||
use_mpnn: bool = False,
|
||
# CheMeleon
|
||
use_chemeleon: bool = False,
|
||
chemeleon_cache: str = "data/processed/chemeleon_embeddings.npz",
|
||
use_unimol: bool = False,
|
||
unimol_cache: str = "data/processed/unimol_embeddings.npz",
|
||
use_moleculestm: bool = False,
|
||
moleculestm_cache: str = "data/processed/moleculestm_embeddings.npz",
|
||
use_mole: bool = False,
|
||
mole_cache: str = "data/processed/mole_embeddings.npz",
|
||
n_repeats: int = 1,
|
||
repeat_seed_step: int = 1000,
|
||
# MoE（消融开关）
|
||
use_moe: bool = False,
|
||
moe_n_experts: int = 4,
|
||
moe_top_k: int = 2,
|
||
moe_expert_hidden_mult: int = 2,
|
||
moe_jitter_noise: float = 0.0,
|
||
# Set Transformer
|
||
set_transformer_block: str = "sab",
|
||
# 回归旁路（消融开关）
|
||
reg_bypass: str = "on",
|
||
# LLM（消融开关）
|
||
use_llm: bool = False,
|
||
use_rag: bool = False,
|
||
rag_top_k: int = 4,
|
||
use_retrieval: bool = False,
|
||
retrieval_source: str = "internal",
|
||
llm_model_path: str = DEFAULT_MOLT5_PATH,
|
||
llm_freeze: bool = True,
|
||
llm_use_lora: bool = False,
|
||
llm_use_qlora: bool = False,
|
||
use_soft_prompt: bool = False,
|
||
llm_lora_r: int = 8,
|
||
llm_lora_alpha: int = 16,
|
||
llm_lora_dropout: float = 0.05,
|
||
# 并行
|
||
parallel: bool = False,
|
||
# 设备
|
||
device: str = "cuda" if torch.cuda.is_available() else "cpu",
|
||
):
|
||
"""
|
||
嵌套交叉验证 + Optuna 超参调优。
|
||
|
||
外层 5-fold（20% test / 80% train），内层 3-fold Optuna 调参。
|
||
外层训练不使用 early-stopping，epoch 数使用内层 best trial 的 epoch_mean。
|
||
|
||
使用 --init-from-pretrain 从预训练 checkpoint 初始化模型权重。
|
||
使用 --parallel 同时运行所有外层 fold（需要足够 GPU 显存）。
|
||
"""
|
||
if optuna is None:
|
||
logger.error("Optuna not installed. Run: pip install optuna")
|
||
raise typer.Exit(1)
|
||
|
||
logger.info(f"Using device: {device}")
|
||
device = torch.device(device)
|
||
set_global_seed(seed)
|
||
|
||
llm_kwargs = dict(
|
||
reg_bypass=reg_bypass,
|
||
use_rag=use_rag,
|
||
rag_top_k=rag_top_k,
|
||
use_llm=use_llm,
|
||
llm_model_path=llm_model_path,
|
||
llm_freeze=llm_freeze,
|
||
llm_use_lora=llm_use_lora,
|
||
llm_use_qlora=llm_use_qlora,
|
||
use_soft_prompt=use_soft_prompt,
|
||
llm_lora_r=llm_lora_r,
|
||
llm_lora_alpha=llm_lora_alpha,
|
||
llm_lora_dropout=llm_lora_dropout,
|
||
)
|
||
|
||
# 加载预训练权重（如果指定）
|
||
pretrain_state_dict = None
|
||
pretrain_config = None
|
||
if init_from_pretrain is not None:
|
||
if init_from_pretrain.exists():
|
||
logger.info(f"Loading pretrain weights from {init_from_pretrain}")
|
||
checkpoint = torch.load(init_from_pretrain, map_location="cpu", weights_only=False)
|
||
pretrain_state_dict = checkpoint["model_state_dict"]
|
||
pretrain_config = checkpoint.get("config", {})
|
||
logger.success(f"Loaded pretrain checkpoint (d_model={pretrain_config.get('d_model')})")
|
||
else:
|
||
logger.warning(f"Pretrain checkpoint not found: {init_from_pretrain}, skipping")
|
||
|
||
# 创建输出目录（带时间戳；--resume-dir 指定则复用，支持断点续跑）
|
||
if resume_dir is not None:
|
||
run_dir = Path(resume_dir)
|
||
run_dir.mkdir(parents=True, exist_ok=True)
|
||
logger.info(f"[RESUME] 复用已有运行目录: {run_dir}")
|
||
else:
|
||
run_name = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||
run_dir = output_dir / run_name
|
||
run_dir.mkdir(parents=True, exist_ok=True)
|
||
logger.info(f"Output directory: {run_dir}")
|
||
|
||
# 加载数据
|
||
logger.info(f"Loading data from {input_path}")
|
||
df = pd.read_csv(input_path)
|
||
logger.info(f"Loaded {len(df)} samples")
|
||
|
||
# 处理数据
|
||
logger.info("Processing dataframe...")
|
||
df = process_dataframe(df)
|
||
|
||
# 构建复合分层标签
|
||
logger.info("Building composite strata...")
|
||
strata, strata_info = build_composite_strata(df, min_stratum_count)
|
||
|
||
# 保存 strata 信息
|
||
with open(run_dir / "strata_info.json", "w") as f:
|
||
json.dump(strata_info, f, indent=2, default=str)
|
||
|
||
# 创建完整数据集（仅用于获取样本数做 split）
|
||
n_samples = len(LNPDataset(df))
|
||
|
||
# 外层 CV split
|
||
outer_cv = StratifiedKFold(
|
||
n_splits=n_outer_folds, shuffle=True, random_state=seed
|
||
)
|
||
|
||
device_str = str(device)
|
||
fold_args = []
|
||
for outer_fold, (outer_train_idx, outer_test_idx) in enumerate(
|
||
outer_cv.split(np.arange(n_samples), strata)
|
||
):
|
||
fold_args.append(dict(
|
||
outer_fold=outer_fold,
|
||
outer_train_idx=outer_train_idx,
|
||
outer_test_idx=outer_test_idx,
|
||
df=df,
|
||
strata=strata,
|
||
fold_dir=run_dir / f"outer_fold_{outer_fold}",
|
||
n_trials=n_trials,
|
||
epochs_per_trial=epochs_per_trial,
|
||
inner_patience=inner_patience,
|
||
batch_size=batch_size,
|
||
n_inner_folds=n_inner_folds,
|
||
use_mpnn=use_mpnn,
|
||
chemeleon_cache=(chemeleon_cache if use_chemeleon else None),
|
||
unimol_cache=(unimol_cache if use_unimol else None),
|
||
moleculestm_cache=(moleculestm_cache if use_moleculestm else None),
|
||
mole_cache=(mole_cache if use_mole else None),
|
||
seed=seed,
|
||
pretrain_state_dict=pretrain_state_dict,
|
||
pretrain_config=pretrain_config,
|
||
load_delivery_head=load_delivery_head,
|
||
device_str=device_str,
|
||
use_retrieval=use_retrieval,
|
||
retrieval_source=retrieval_source,
|
||
use_moe=use_moe,
|
||
moe_n_experts=moe_n_experts,
|
||
moe_top_k=moe_top_k,
|
||
moe_expert_hidden_mult=moe_expert_hidden_mult,
|
||
moe_jitter_noise=moe_jitter_noise,
|
||
set_transformer_block=set_transformer_block,
|
||
llm_kwargs=llm_kwargs,
|
||
))
|
||
|
||
if parallel:
|
||
import multiprocessing as mp
|
||
from concurrent.futures import ProcessPoolExecutor
|
||
|
||
ctx = mp.get_context("spawn")
|
||
logger.info(f"Running {n_outer_folds} outer folds in PARALLEL (spawn)")
|
||
with ProcessPoolExecutor(max_workers=n_outer_folds, mp_context=ctx) as executor:
|
||
futures = [executor.submit(_run_single_outer_fold, **args) for args in fold_args]
|
||
outer_results = [f.result() for f in futures]
|
||
outer_results.sort(key=lambda r: r["fold"])
|
||
else:
|
||
logger.info(f"Running {n_outer_folds} outer folds SEQUENTIALLY")
|
||
outer_results = []
|
||
per_fold_params = {}
|
||
for args in fold_args:
|
||
result = _run_single_outer_fold(**args)
|
||
outer_results.append(result)
|
||
per_fold_params[result["fold"]] = (result["best_params"], result["epoch_mean"])
|
||
|
||
for rep in range(1, n_repeats):
|
||
logger.info(f"\n===== Repeat {rep}/{n_repeats - 1} (reuse hyperparams) =====")
|
||
for args in fold_args:
|
||
bp, em = per_fold_params[args["outer_fold"]]
|
||
rep_args = dict(args)
|
||
rep_args["seed"] = seed + rep * repeat_seed_step
|
||
rep_args["fold_dir"] = run_dir / f"repeat_{rep}" / f"outer_fold_{args['outer_fold']}"
|
||
rep_args["precomputed_best_params"] = bp
|
||
rep_args["precomputed_epoch_mean"] = em
|
||
outer_results.append(_run_single_outer_fold(**rep_args))
|
||
|
||
# 汇总结果
|
||
logger.info("\n" + "=" * 60)
|
||
logger.info("NESTED CV COMPLETE")
|
||
logger.info("=" * 60)
|
||
|
||
# 计算汇总统计
|
||
summary = {"fold_results": outer_results}
|
||
|
||
# 对每个任务计算均值和标准差
|
||
tasks_with_metrics = {}
|
||
for result in outer_results:
|
||
for task, metrics in result["test_metrics"].items():
|
||
if task not in tasks_with_metrics:
|
||
tasks_with_metrics[task] = {k: [] for k in metrics.keys() if k != "n_samples"}
|
||
for k, v in metrics.items():
|
||
if k != "n_samples":
|
||
tasks_with_metrics[task][k].append(v)
|
||
|
||
summary["summary_stats"] = {}
|
||
for task, metrics_dict in tasks_with_metrics.items():
|
||
summary["summary_stats"][task] = {}
|
||
for metric_name, values in metrics_dict.items():
|
||
summary["summary_stats"][task][f"{metric_name}_mean"] = float(np.mean(values))
|
||
summary["summary_stats"][task][f"{metric_name}_std"] = float(np.std(values))
|
||
|
||
# 打印汇总
|
||
logger.info("\n[Summary Statistics]")
|
||
for task, stats in summary["summary_stats"].items():
|
||
if "rmse_mean" in stats:
|
||
logger.info(
|
||
f" {task}: RMSE={stats['rmse_mean']:.4f}±{stats['rmse_std']:.4f}, "
|
||
f"R²={stats['r2_mean']:.4f}±{stats['r2_std']:.4f}"
|
||
)
|
||
elif "accuracy_mean" in stats:
|
||
logger.info(
|
||
f" {task}: Acc={stats['accuracy_mean']:.4f}±{stats['accuracy_std']:.4f}, "
|
||
f"F1={stats['f1_mean']:.4f}±{stats['f1_std']:.4f}"
|
||
)
|
||
elif "kl_divergence_mean" in stats:
|
||
logger.info(
|
||
f" {task}: KL={stats['kl_divergence_mean']:.4f}±{stats['kl_divergence_std']:.4f}, "
|
||
f"JS={stats['js_divergence_mean']:.4f}±{stats['js_divergence_std']:.4f}"
|
||
)
|
||
|
||
# 保存汇总
|
||
with open(run_dir / "summary.json", "w") as f:
|
||
json.dump(summary, f, indent=2)
|
||
|
||
logger.success(f"\nAll results saved to {run_dir}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
app()
|
||
|