diff --git a/lnp_ml/modeling/final_train_optuna_cv.py b/lnp_ml/modeling/final_train_optuna_cv.py index ede7852..7a94fba 100644 --- a/lnp_ml/modeling/final_train_optuna_cv.py +++ b/lnp_ml/modeling/final_train_optuna_cv.py @@ -491,6 +491,8 @@ def main( n_trials: int = 20, epochs_per_trial: int = 30, patience: int = 10, + fixed_params_json: Optional[Path] = None, + fixed_epoch_mean: Optional[int] = None, # 训练参数 batch_size: int = 32, # 最终训练参数 @@ -613,75 +615,86 @@ def main( # 预热 RDKit 缓存（在整个训练流程中共享） rdkit_cache = warmup_rdkit_cache(full_dataset.smiles) - # 运行 Optuna 调参 - logger.info(f"\nRunning {n_folds}-fold Optuna with {n_trials} trials...") - study_path = output_dir / "optuna_study.sqlite3" - - best_params, epoch_mean, study = run_optuna_cv( - full_dataset=full_dataset, - strata=strata, - device=device, - n_trials=n_trials, - epochs_per_trial=epochs_per_trial, - patience=patience, - batch_size=batch_size, - n_folds=n_folds, - use_mpnn=use_mpnn, - chemeleon_cache=(chemeleon_cache if use_chemeleon else None), - unimol_cache=(unimol_cache if use_unimol else None), - seed=seed, - 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, - use_retrieval=use_retrieval, - freeze_backbone_epochs=freeze_backbone_epochs, - ) - + if fixed_params_json is not None: + if fixed_epoch_mean is None: + raise typer.BadParameter("--fixed-params-json 必须配合 --fixed-epoch-mean 使用") + logger.info(f"Skipping Optuna, loading fixed params from {fixed_params_json}") + with open(fixed_params_json) as f: + best_params = json.load(f) + epoch_mean = fixed_epoch_mean + study = None + logger.info(f"Fixed params: {best_params}") + logger.info(f"Fixed epoch_mean: {epoch_mean}") + else: + logger.info(f"\nRunning {n_folds}-fold Optuna with {n_trials} trials...") + study_path = output_dir / "optuna_study.sqlite3" + + best_params, epoch_mean, study = run_optuna_cv( + full_dataset=full_dataset, + strata=strata, + device=device, + n_trials=n_trials, + epochs_per_trial=epochs_per_trial, + patience=patience, + batch_size=batch_size, + n_folds=n_folds, + use_mpnn=use_mpnn, + chemeleon_cache=(chemeleon_cache if use_chemeleon else None), + unimol_cache=(unimol_cache if use_unimol else None), + seed=seed, + 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, + use_retrieval=use_retrieval, + freeze_backbone_epochs=freeze_backbone_epochs, + ) + # 保存最佳参数 with open(output_dir / "best_params.json", "w") as f: json.dump(best_params, f, indent=2) - + with open(output_dir / "epoch_mean.json", "w") as f: json.dump({"epoch_mean": epoch_mean}, f) - + # 保存 Optuna 试验历史 - trials_history = [] - for trial in study.trials: - trials_history.append({ - "number": trial.number, - "value": trial.value, - "params": trial.params, - "user_attrs": trial.user_attrs, - "state": str(trial.state), - }) - - with open(output_dir / "optuna_trials.json", "w") as f: - json.dump(trials_history, f, indent=2) - + if study is not None: + trials_history = [] + for trial in study.trials: + trials_history.append({ + "number": trial.number, + "value": trial.value, + "params": trial.params, + "user_attrs": trial.user_attrs, + "state": str(trial.state), + }) + + with open(output_dir / "optuna_trials.json", "w") as f: + json.dump(trials_history, f, indent=2) + # 全量数据训练 logger.info(f"\n{'='*60}") logger.info("FINAL TRAINING ON FULL DATA") logger.info(f"{'='*60}") logger.info(f"Using best params with epochs={epoch_mean}") logger.info(f"SWA: {use_swa}") - + # 创建全量 DataLoader full_loader = DataLoader( full_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn ) - + # 计算类权重 class_weights = compute_class_weights_from_loader(full_loader) - + # 保存类权重信息 class_weights_info = { "pdi": class_weights.pdi.tolist() if class_weights.pdi is not None else None, @@ -690,9 +703,7 @@ def main( } with open(output_dir / "class_weights.json", "w") as f: json.dump(class_weights_info, f, indent=2) - - # 架构超参一次性解析：Optuna 采样值优先于 CLI 默认值。 - # 解析结果必须同时喂给 create_model 和 config，否则 load_model 回读时结构对不上。 + arch = { "moe_n_experts": best_params.get("moe_n_experts", moe_n_experts), "moe_top_k": best_params.get("moe_top_k", moe_top_k), diff --git a/models/final/best_params.json b/models/final/best_params.json index df17e17..7395ce5 100644 --- a/models/final/best_params.json +++ b/models/final/best_params.json @@ -3,7 +3,7 @@ "lr": 0.000980484886752915, "weight_decay": 0.00014367509113739864, "backbone_lr_ratio": 0.13093310334961422, - "moe_n_experts": 2, + "moe_n_experts": 4, "moe_top_k": 2, "moe_expert_hidden_mult": 2, "llm_lora_r": 8, diff --git a/models/final/history.json b/models/final/history.json index 2d1eabb..dd5082f 100644 --- a/models/final/history.json +++ b/models/final/history.json @@ -1,126 +1,126 @@ { "train": [ { - "loss": 5.662022154286222, - "loss_size": 0.9464401806581695, - "loss_pdi": 0.6576592483610477, - "loss_ee": 1.0441054870497506, - "loss_delivery": 1.0691580114499577, - "loss_biodist": 0.7718977461446006, - "loss_toxic": 0.47160032813279135, - "loss_moe_lb": 1.9999999887538407, - "loss_aux_moe": 1.096376850709038, - "loss_aux_llm": 1.2355621509113401 + "loss": 5.694288951046062, + "loss_size": 1.015234741259296, + "loss_pdi": 0.6592788083373375, + "loss_ee": 1.027098409409793, + "loss_delivery": 1.0114885227016683, + "loss_biodist": 0.7630386647750746, + "loss_toxic": 0.5493527664890829, + "loss_moe_lb": 2.045894391131851, + "loss_aux_moe": 1.0363417644545716, + "loss_aux_llm": 1.1829107824080396 }, { - "loss": 4.549718969273117, - "loss_size": 0.9265210998227011, - "loss_pdi": 0.6320651516599475, - "loss_ee": 0.9564013177493833, - "loss_delivery": 0.8463795804682205, - "loss_biodist": 0.48631338606465535, - "loss_toxic": 0.25031149114991696, - "loss_moe_lb": 1.9999999910030726, - "loss_aux_moe": 0.8730144033928946, - "loss_aux_llm": 1.0696029401612732 + "loss": 4.587394336484513, + "loss_size": 0.9309988603839334, + "loss_pdi": 0.6420835014784111, + "loss_ee": 0.9480358497151788, + "loss_delivery": 0.809673407519201, + "loss_biodist": 0.5112692536610477, + "loss_toxic": 0.30279909285171974, + "loss_moe_lb": 2.0554285611746446, + "loss_aux_moe": 0.9779460283242306, + "loss_aux_llm": 0.9288384830754883 }, { - "loss": 4.183341939494295, - "loss_size": 0.8995583306927726, - "loss_pdi": 0.6350391439671786, - "loss_ee": 0.9038635166186206, - "loss_delivery": 0.8323076782080362, - "loss_biodist": 0.44391101247297143, - "loss_toxic": 0.1671132598740031, - "loss_moe_lb": 1.9999999865046088, - "loss_aux_moe": 0.8456886384003567, - "loss_aux_llm": 1.0173823231796049 + "loss": 4.160386994199933, + "loss_size": 0.9387357568122306, + "loss_pdi": 0.5914091456611201, + "loss_ee": 0.9026568987459507, + "loss_delivery": 0.84944020264413, + "loss_biodist": 0.3876211567307418, + "loss_toxic": 0.18268741205563102, + "loss_moe_lb": 2.1000781711542382, + "loss_aux_moe": 0.9373761130461715, + "loss_aux_llm": 0.9303386138295228 }, { - "loss": 3.9456988685535936, - "loss_size": 0.925582969104344, - "loss_pdi": 0.586557297211773, - "loss_ee": 0.867810919599713, - "loss_delivery": 0.7656616399521535, - "loss_biodist": 0.4162127935099152, - "loss_toxic": 0.2137195658123226, - "loss_moe_lb": 1.9999999887538407, - "loss_aux_moe": 0.7506198146784643, - "loss_aux_llm": 0.9585366686981804 + "loss": 3.809376889804624, + "loss_size": 0.888136670766574, + "loss_pdi": 0.5791616782827197, + "loss_ee": 0.8553679860987753, + "loss_delivery": 0.8118948299648627, + "loss_biodist": 0.30093744889182866, + "loss_toxic": 0.15547859341890183, + "loss_moe_lb": 2.102432388179707, + "loss_aux_moe": 1.1255526001442153, + "loss_aux_llm": 0.9655606655298539 }, { - "loss": 3.4442428957741216, - "loss_size": 0.8259714532573268, - "loss_pdi": 0.5488103347004585, - "loss_ee": 0.8086997242468708, - "loss_delivery": 0.7543281304808158, - "loss_biodist": 0.3298270061330975, - "loss_toxic": 0.09599270570566351, - "loss_moe_lb": 1.9999999955015362, - "loss_aux_moe": 0.8137425728282839, - "loss_aux_llm": 1.044384686873769 + "loss": 3.1629226522625618, + "loss_size": 0.7107668774870207, + "loss_pdi": 0.5377992931401955, + "loss_ee": 0.7858308424364846, + "loss_delivery": 0.7225946192190332, + "loss_biodist": 0.23169103110173964, + "loss_toxic": 0.10101612598441963, + "loss_moe_lb": 2.1294895545491634, + "loss_aux_moe": 0.8120556049189478, + "loss_aux_llm": 0.8958881971127582 }, { - "loss": 3.108477974837681, - "loss_size": 0.7671496489981435, - "loss_pdi": 0.5366503028374798, - "loss_ee": 0.7869719379353073, - "loss_delivery": 0.685029683248052, - "loss_biodist": 0.2828179093183212, - "loss_toxic": 0.08221106674908749, - "loss_moe_lb": 1.9999999910030726, - "loss_aux_moe": 0.7168521201413758, - "loss_aux_llm": 0.9446638939234445 + "loss": 2.988117564399287, + "loss_size": 0.7456601148225227, + "loss_pdi": 0.5200111323370123, + "loss_ee": 0.7696504339856921, + "loss_delivery": 0.702117268874679, + "loss_biodist": 0.19998157122787438, + "loss_toxic": 0.08327638518992143, + "loss_moe_lb": 2.1596880116552675, + "loss_aux_moe": 0.763107772424536, + "loss_aux_llm": 0.8947458216604197 }, { - "loss": 2.983976681277437, - "loss_size": 0.7495474740159962, - "loss_pdi": 0.5227343659355955, - "loss_ee": 0.7467741780685928, - "loss_delivery": 0.6811002301368511, - "loss_biodist": 0.24798926155803339, - "loss_toxic": 0.11302385037876929, - "loss_moe_lb": 1.9999999955015362 + "loss": 2.7413068245042047, + "loss_size": 0.6423791460692883, + "loss_pdi": 0.5119175066081982, + "loss_ee": 0.7589052460103665, + "loss_delivery": 0.6498366547924168, + "loss_biodist": 0.1679054581612911, + "loss_toxic": 0.09372632653026332, + "loss_moe_lb": 2.148664339533392 }, { - "loss": 3.014890724757932, - "loss_size": 0.7314784642098084, - "loss_pdi": 0.48637263083233023, - "loss_ee": 0.769600696158859, - "loss_delivery": 0.8273954027912246, - "loss_biodist": 0.21140338072799286, - "loss_toxic": 0.08676587497249338, - "loss_moe_lb": 1.9999999955015362 + "loss": 2.7408997472727075, + "loss_size": 0.6384303912801562, + "loss_pdi": 0.49059138151834597, + "loss_ee": 0.7721629592607606, + "loss_delivery": 0.7107182373556326, + "loss_biodist": 0.15433137111787526, + "loss_toxic": 0.0748163205291417, + "loss_moe_lb": 2.154493790752483 }, { - "loss": 2.774192355713754, - "loss_size": 0.6206410539881239, - "loss_pdi": 0.4826473706173447, - "loss_ee": 0.7343562110415045, - "loss_delivery": 0.6946425725758638, - "loss_biodist": 0.19993741290186937, - "loss_toxic": 0.1272664367724298, - "loss_moe_lb": 1.9999999932523043 + "loss": 2.610751993251297, + "loss_size": 0.6624566971693399, + "loss_pdi": 0.4767126595636584, + "loss_ee": 0.7076389159796372, + "loss_delivery": 0.636365042331646, + "loss_biodist": 0.15519743205382014, + "loss_toxic": 0.06698143800538262, + "loss_moe_lb": 2.1544579132547916 }, { - "loss": 2.683339330385316, - "loss_size": 0.5581827477885867, - "loss_pdi": 0.4963476042140205, - "loss_ee": 0.7145765797709519, - "loss_delivery": 0.6466561276817097, - "loss_biodist": 0.1851363978436533, - "loss_toxic": 0.1616253986507698, - "loss_moe_lb": 1.9999999932523043 + "loss": 2.581533803130096, + "loss_size": 0.6987560360499148, + "loss_pdi": 0.4562466161431007, + "loss_ee": 0.7065791735109294, + "loss_delivery": 0.6151787998541346, + "loss_biodist": 0.1382373322223155, + "loss_toxic": 0.06480917761778845, + "loss_moe_lb": 2.137455944745046 }, { - "loss": 2.5718357652988075, - "loss_size": 0.6112437141391466, - "loss_pdi": 0.4928069238392812, - "loss_ee": 0.7109674007262824, - "loss_delivery": 0.6155119884829476, - "loss_biodist": 0.18839857628885306, - "loss_toxic": 0.05047845465344166, - "loss_moe_lb": 1.9999999842553768 + "loss": 2.3695818658144967, + "loss_size": 0.5367685263308714, + "loss_pdi": 0.47306264318385216, + "loss_ee": 0.6626323349070999, + "loss_delivery": 0.5844551013817765, + "loss_biodist": 0.13707212931564394, + "loss_toxic": 0.06755097333027735, + "loss_moe_lb": 2.1183437378901355 } ], "val": [] diff --git a/models/final/loss_curves.png b/models/final/loss_curves.png index 59f3b98..a031191 100644 Binary files a/models/final/loss_curves.png and b/models/final/loss_curves.png differ diff --git a/models/final/model.pt b/models/final/model.pt index 9ea1e33..9e0d922 100644 --- a/models/final/model.pt +++ b/models/final/model.pt @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:b96d3867be53a3e23abdf291baeec4a9d610da58dc790220a345e7012187aaa0 -size 53008646 +oid sha256:5b1c826c11338f313eea643f839fbc27c4c1e0f0f19d58d672d405cc731caa2a +size 58261734 diff --git a/scripts/hparams_sparse_moe.json b/scripts/hparams_sparse_moe.json new file mode 100644 index 0000000..7395ce5 --- /dev/null +++ b/scripts/hparams_sparse_moe.json @@ -0,0 +1,16 @@ +{ + "dropout": 0.19323529834582587, + "lr": 0.000980484886752915, + "weight_decay": 0.00014367509113739864, + "backbone_lr_ratio": 0.13093310334961422, + "moe_n_experts": 4, + "moe_top_k": 2, + "moe_expert_hidden_mult": 2, + "llm_lora_r": 8, + "d_model": 256, + "num_heads": 8, + "n_attn_layers": 4, + "fusion_strategy": "attention", + "head_hidden_dim": 128, + "set_transformer_block": "sab" +} \ No newline at end of file diff --git a/scripts_run/run_final_full.sh b/scripts_run/run_final_full.sh index a0f30de..6376256 100644 --- a/scripts_run/run_final_full.sh +++ b/scripts_run/run_final_full.sh @@ -11,6 +11,8 @@ N_TRIALS=${N_TRIALS:-20} EPOCHS=${EPOCHS:-20} # 与 nested CV 的 EPOCHS 保持一致 PATIENCE=${PATIENCE:-5} # 与 nested CV 的 PATIENCE 保持一致 N_FOLDS=${N_FOLDS:-3} +FIXED_PARAMS=${FIXED_PARAMS:-} # 非空则跳过 Optuna，直接用这份超参 +FIXED_EPOCH_MEAN=${FIXED_EPOCH_MEAN:-11} BATCH=${BATCH:-8} REG_BYPASS=${REG_BYPASS:-off} FREEZE=${FREEZE:-3} @@ -63,6 +65,12 @@ if [ -n "${PRETRAIN}" ] && [ -f "${PRETRAIN}" ]; then elif [ -n "${PRETRAIN}" ]; then echo "[$(date '+%F %T')] 警告：${PRETRAIN} 不存在，跳过预训练初始化" >>"${LOG}" fi +if [ -n "${FIXED_PARAMS}" ]; then + if [ ! -f "${FIXED_PARAMS}" ]; then + echo "[$(date '+%F %T')] 错误：${FIXED_PARAMS} 不存在" >>"${LOG}"; exit 1 + fi + EXTRA+=(--fixed-params-json "${FIXED_PARAMS}" --fixed-epoch-mean "${FIXED_EPOCH_MEAN}") +fi for attempt in $(seq 1 "${MAX_RETRY}"); do wait_free