src/wjad/train/runner_local.py

"""本地训练入口。

支持两种模式：
  - ``--tiny`` : 用 1 个 clip / 极少步数验证训练循环可跑通；
  - 否则默认按 configs/default.yaml 训练（需要数据集解压完成）。
"""

from __future__ import annotations

import argparse
import logging
import os
import sys
from pathlib import Path

import torch
import yaml
from torch.utils.data import DataLoader

from ..data.cosmos_dataset import CosmosDriveDreamsDataset, build_clip_index, collate_samples
from ..model import E2EAVModel
from .trainer import Trainer, TrainerConfig


def _load_config(path: str) -> dict:
    with open(path, "r", encoding="utf-8") as f:
        return yaml.safe_load(f)


def _deep_merge(base: dict, override: dict) -> dict:
    out = dict(base)
    for k, v in override.items():
        if k in out and isinstance(out[k], dict) and isinstance(v, dict):
            out[k] = _deep_merge(out[k], v)
        else:
            out[k] = v
    return out


def _make_model_from_cfg(cfg: dict, dinov3_path: str) -> E2EAVModel:
    """根据配置创建模型。"""
    return E2EAVModel(
        dinov3_path=dinov3_path,
        backbone_dim=cfg["backbone"]["hidden_size"],
        num_heads=cfg["backbone"]["num_heads"],
        num_dense_layers=cfg["backbone"]["num_dense_layers"],
        num_moe_layers=cfg["backbone"]["num_moe_layers"],
        num_routed_experts=cfg["moe"]["num_routed_experts"],
        num_shared_experts=cfg["moe"]["num_shared_experts"],
        topk_experts=cfg["moe"]["topk"],
        ffn_mult=cfg["backbone"]["ffn_mult"],
        num_history_frames=cfg["input"]["num_history_frames"],
        num_detection_tokens=cfg["tokens"]["num_detection"],
        num_control_tokens=cfg["tokens"]["num_control"],
        num_ego_tokens=cfg["tokens"]["num_ego"],
        num_extra_tokens=cfg["tokens"]["num_extra"],
        image_h=cfg["input"]["image_height"],
        image_w=cfg["input"]["image_width"],
        patch_size=cfg["dinov3"]["patch_size"],
        num_classes=cfg["det_traj_head"]["num_classes"],
        traj_horizon=cfg["det_traj_head"]["traj_horizon"],
        det_head_hidden=cfg["det_traj_head"]["hidden_size"],
        ctrl_head_hidden=cfg["control_head"]["hidden_size"],
        calib_dim=cfg["calibration"]["hidden_size"],
        calib_num_query=cfg["calibration"]["num_query_tokens"],
        calib_num_blocks=cfg["calibration"]["num_blocks"],
        calib_num_self_per_block=cfg["calibration"]["num_self_attn_per_block"],
        calib_num_heads=cfg["calibration"]["num_heads"],
        calib_residual_range=cfg["calibration"]["residual_range"],
        calib_intr_dim=cfg["calibration"]["intr_vec_dim"],
        freeze_dinov3=cfg["dinov3"]["freeze_in_stage1"],
        attn_implementation=cfg["dinov3"]["attn_implementation"],
    )


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default=str(Path(__file__).resolve().parents[3] / "configs" / "default.yaml"))
    parser.add_argument(
        "--config_overrides",
        default=None,
        help="可选第二份 YAML，与 --config 深度合并（如 configs/jobs_overrides.yaml）",
    )
    parser.add_argument("--data_root", default=None, help="覆盖 config 中的 data.root")
    parser.add_argument("--dinov3_path", default=str(Path(__file__).resolve().parents[3] / "dinov3-vitb16-pretrain-lvd1689m"))
    parser.add_argument("--tiny", action="store_true", help="用极少样本验证训练循环")
    parser.add_argument(
        "--max_steps",
        type=int,
        default=None,
        help="默认：tiny=50；全量训练=None 使用配置文件 total_steps",
    )
    parser.add_argument("--device", default="cpu")
    parser.add_argument("--output_dir", default=None, help="检查点目录；也可用环境变量 WJAD_OUTPUT_DIR")
    parser.add_argument("--hub_repo", default=None, help="Hub model repo（推送 checkpoint）；或 WJAD_HUB_REPO")
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
    cfg = _load_config(args.config)
    if args.config_overrides:
        cfg = _deep_merge(cfg, _load_config(args.config_overrides))
    data_root = args.data_root or cfg["data"]["root"]

    samples = build_clip_index(
        data_root,
        weathers=cfg["data"]["weather"],
        camera_name=cfg["input"]["camera_name"],
    )
    if args.tiny:
        samples = samples[:8] or samples
    print(f"[runner_local] 找到 {len(samples)} 个样本于 {data_root}")
    if not samples:
        print("[runner_local] 没有数据。请先运行 scripts/download_data.py 准备数据集。")
        sys.exit(0)

    dataset = CosmosDriveDreamsDataset(
        data_root=data_root,
        samples=samples,
        camera_name=cfg["input"]["camera_name"],
        image_h=cfg["input"]["image_height"],
        image_w=cfg["input"]["image_width"],
        num_history=cfg["input"]["num_history_frames"],
        future_horizon=cfg["input"]["num_future_frames"],
        max_distance_m=cfg["detection"]["max_distance_m"],
        occlusion_tol=cfg["detection"]["occlusion_depth_tolerance"],
    )
    bs = cfg["train"]["batch_size"]
    if os.environ.get("WJAD_BATCH_SIZE"):
        bs = int(os.environ["WJAD_BATCH_SIZE"])

    deploy = cfg.get("deploy") or {}
    output_dir = args.output_dir or os.environ.get("WJAD_OUTPUT_DIR")
    hub_repo = args.hub_repo or os.environ.get("WJAD_HUB_REPO") or deploy.get("hf_weights_repo")

    loader = DataLoader(
        dataset,
        batch_size=bs,
        shuffle=True,
        num_workers=cfg["data"]["num_workers"] if not args.tiny else 0,
        collate_fn=collate_samples,
        pin_memory=cfg["data"]["pin_memory"],
    )

    model = _make_model_from_cfg(cfg, args.dinov3_path)
    if cfg.get("gradient_checkpointing", False):
        model.backbone.set_gradient_checkpointing(True)

    tcfg = TrainerConfig(
        total_steps=cfg["train"]["total_steps"],
        warmup_steps=cfg["train"]["warmup_steps"],
        base_lr=cfg["train"]["base_lr"],
        min_lr=cfg["train"]["min_lr"],
        weight_decay=cfg["train"]["weight_decay"],
        grad_clip=cfg["train"]["grad_clip"],
        log_interval=cfg["train"]["log_interval"],
        ckpt_interval=cfg["train"]["ckpt_interval"],
        stage1_steps=cfg["train"]["stage1_steps"],
        stage1_perturb_start=cfg["train"]["stage1_perturb_start"],
        grad_monitor_threshold=cfg["train"]["grad_monitor_threshold"],
        moe_load_balance_weight=cfg["moe"]["load_balance_weight"],
        moe_boundary_weight=cfg["moe"]["boundary_weight"],
        router_temp_init=cfg["moe"]["router_temperature_init"],
        router_temp_final=cfg["moe"]["router_temperature_final"],
        loss_giou_weight=cfg["loss"]["giou_weight"],
        loss_calib_weight=cfg["loss"]["calib_weight"],
        enable_gradnorm=cfg["multitask"]["enable_gradnorm"],
        enable_pcgrad=cfg["multitask"]["enable_pcgrad"],
        mixed_precision=cfg["mixed_precision"],
        grad_accum_steps=cfg["train"]["grad_accum_steps"],
        dinov3_lr_mult_stage2=cfg["dinov3"]["finetune_lr_ratio"],
        backbone_lr_mult=cfg["train"]["param_groups"]["backbone_lr_mult"],
        calibration_lr_mult=cfg["train"]["param_groups"]["calibration_lr_mult"],
        head_lr_mult=cfg["train"]["param_groups"]["head_lr_mult"],
        gate_lr_mult=cfg["train"]["param_groups"]["gate_lr_mult"],
        output_dir=output_dir,
        hub_repo_id=hub_repo,
    )
    trainer = Trainer(model, tcfg, num_classes=cfg["det_traj_head"]["num_classes"], device=args.device)
    if args.max_steps is not None:
        max_steps = args.max_steps
    else:
        max_steps = 50 if args.tiny else None
    trainer.fit(loader, max_steps=max_steps)


if __name__ == "__main__":
    main()