File size: 5,774 Bytes
be0e2ec | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | import json
import copy
import shutil
from pathlib import Path
import optuna
from optuna.samplers import TPESampler
from optuna.pruners import MedianPruner
from transformers import TrainerCallback
from src.training.train_model import (
load_yaml_config,
resolve_runtime_config,
train_once,
)
class OptunaPruningCallback(TrainerCallback):
"""Report dev-metric mỗi lần eval -> MedianPruner cắt sớm trial tệ (Phase 02).
Pruner cũ vô dụng vì objective chỉ trả best cuối, KHÔNG report intermediate."""
def __init__(self, trial, metric_name: str = "eval_macro_f1"):
self.trial = trial
self.metric_name = metric_name
self.step = 0
def on_evaluate(self, args, state, control, metrics=None, **kwargs):
if not metrics or self.metric_name not in metrics:
return
self.trial.report(metrics[self.metric_name], step=self.step)
self.step += 1
if self.trial.should_prune():
raise optuna.TrialPruned()
def objective(trial: optuna.Trial, base_config: dict) -> float:
config = copy.deepcopy(base_config)
model_type = config["model"].get("type", "encoder")
if model_type == "decoder_lora":
# LEAN: mỗi trial = train LLM 7B (~20-30 phút) -> không sweep batch/max_length.
config["training"]["learning_rate"] = trial.suggest_float("learning_rate", 5e-5, 3e-4, log=True)
config["training"]["epochs"] = trial.suggest_int("epochs", 2, 4)
lora_r = trial.suggest_categorical("lora_r", [8, 16, 32])
config["model"].setdefault("lora", {})
config["model"]["lora"]["r"] = lora_r
config["model"]["lora"]["alpha"] = 2 * lora_r
else:
# encoder (PhoBERT) — search LR/batch/wd/warmup (Phase 02: BỎ epochs, THÊM warmup).
# max_length KHÔNG search: set per-task từ data (topic=128, subst=256) trong config/train.yml.
config["training"]["learning_rate"] = trial.suggest_float("learning_rate", 1e-5, 5e-5, log=True)
# batch search per-task từ config (subst data ít -> [8,16]; topic -> [8,16,32]). Xem config/train.yml.
batch_choices = config.get("tune", {}).get("train_batch_size", [8, 16, 32])
config["training"]["train_batch_size"] = trial.suggest_categorical("train_batch_size", batch_choices)
config["training"]["weight_decay"] = trial.suggest_float("weight_decay", 0.0, 0.10, step=0.0001)
# warmup_ratio: lever ổn định fine-tune (Mosbach/Dodge). KHÔNG tune dropout (chốt 2026-06-13).
config["training"]["warmup_ratio"] = trial.suggest_float("warmup_ratio", 0.0, 0.2)
# epochs KHÔNG tune: cố định = trần config (10) + EarlyStopping quyết best-epoch.
trial_dir = Path(config["paths"]["output_dir"]) / f"trial_{trial.number}"
config["paths"]["output_dir"] = str(trial_dir)
# Không lưu checkpoint trong quá trình tuning để tiết kiệm disk.
# EarlyStoppingCallback vẫn track best_metric qua state nên metric vẫn đúng.
config["training"]["save_strategy"] = "no"
config["training"]["load_best_model_at_end"] = False
print(f"\n{'='*50}")
print(f"TRIAL {trial.number} [{model_type}] — {trial.params}")
print(f"{'='*50}")
metric_name = "eval_" + config["training"].get("metric_for_best_model", "macro_f1")
pruning_cb = OptunaPruningCallback(trial, metric_name=metric_name)
# try/finally: trial bị PRUNE sẽ raise TrialPruned giữa chừng -> phải dọn trial_dir ở finally
# (nếu không, mỗi trial pruned để lại 1 thư mục trial_N rỗng trong outputs/models/<dir>/).
try:
run_results = train_once(config, save_model=False, extra_callbacks=[pruning_cb])
macro_f1 = run_results["trainer"].state.best_metric
print(f"Trial {trial.number} -> Macro-F1: {macro_f1:.4f}")
return macro_f1
finally:
shutil.rmtree(trial_dir, ignore_errors=True)
def main(args_cli=None):
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--config", type=str, default="config/train.yml")
parser.add_argument("--task", type=str, default="topic")
parser.add_argument("--trials", type=int, default=10)
args = parser.parse_args(args_cli)
raw_config = load_yaml_config(Path(args.config))
base_config = resolve_runtime_config(raw_config, task=args.task)
study_name = f"esg_{args.task}_hyperopt"
study = optuna.create_study(
study_name=study_name,
storage=f"sqlite:///{study_name}.db",
direction="maximize",
# n_startup_trials=5: với n_trials=10 -> 5 random thăm dò + 5 TPE khai thác.
# (Mặc định 10 sẽ khiến cả 10 trial là random startup, TPE không bao giờ bật.)
sampler=TPESampler(seed=base_config.get("seed", 42), n_startup_trials=5),
pruner=MedianPruner(n_startup_trials=5, n_warmup_steps=0),
load_if_exists=True,
)
print(f"Optuna tuning: task={args.task}, trials={args.trials}")
study.optimize(
lambda trial: objective(trial, base_config),
n_trials=args.trials,
)
best = study.best_trial
print(f"\nBest trial #{best.number} — Macro-F1: {best.value:.4f}")
print("Params:")
for k, v in best.params.items():
print(f" {k}: {v}")
out_path = Path(base_config["paths"]["output_dir"]) / f"best_params_{args.task}.json"
out_path.parent.mkdir(parents=True, exist_ok=True)
with out_path.open("w", encoding="utf-8") as f:
json.dump(best.params, f, indent=2)
print(f"Saved to {out_path}")
if __name__ == "__main__":
main()
|