train.py

"""
Hi-FL: federated training entry point.

Dataset-agnostic by design -- all dataset-specific details (centers, paths,
folder-naming conventions, class prompts) live in a YAML config under
./configs/. To run on a new dataset, add a new YAML file; no code changes
are required, as long as the dataset follows one of the two layouts
supported by hifl/datasets.py.

Hyperparameter precedence: CLI flag (if explicitly passed) > YAML config
(`model:` / `training:` sections) > hardcoded fallback default. This lets a
config fully specify a reproducible run, while still allowing quick
one-off overrides from the command line.

Usage:
    python train.py --config configs/bypass.yaml
    python train.py --config configs/colorectal.yaml --rounds 30
    python train.py --config configs/multicholec.yaml --margin 0.5
"""
import argparse
import random

import numpy as np
import torch
import yaml
from mmengine.config import Config

import surgvlp
from hifl import HiFLModel, build_loaders, run_federated_training

SEED = 0
torch.manual_seed(SEED)
torch.cuda.manual_seed(SEED)
np.random.seed(SEED)
random.seed(SEED)
torch.backends.cudnn.deterministic = True

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Hardcoded fallbacks, used only if a value is absent from both the CLI and the YAML config.
DEFAULTS = {
    'frac': 0.05, 'rounds': 50,
    'w_tar': 0.3, 'w_kg': 0.1, 'w_indep': 0.1, 'margin': 0.4,
    'lr': 1e-3, 'weight_decay': 1e-4, 'label_smoothing': 0.1, 'grad_clip_norm': 1.0,
    'embed_dim': 768, 'bottleneck_ratio': 4, 'phase_kernel_sizes': [3, 15, 31], 'logit_scale': 100.0,
}


def get_args():
    parser = argparse.ArgumentParser(description="Hi-FL federated training (dataset-agnostic)")
    parser.add_argument('--config', required=True, type=str, help='Path to a dataset YAML config, e.g. configs/bypass.yaml')
    parser.add_argument('--surgvlp_config', default='./tests/config_peskavlp.py', type=str)
    parser.add_argument('--out', default=None, type=str, help='Output checkpoint path for the best global model')

    # All hyperparameters default to None here so we can tell "not passed on
    # the CLI" apart from "explicitly set to a falsy value", and fall back
    # to the YAML config, then to DEFAULTS, in resolve_hparam().
    parser.add_argument('--frac', default=None, type=float, help='Fraction of training videos kept per center')
    parser.add_argument('--rounds', default=None, type=int, help='Number of FedAvg communication rounds')
    # Loss weights
    parser.add_argument('--w_tar', default=None, type=float, help='Weight of the TAR loss')
    parser.add_argument('--w_kg', default=None, type=float, help='Weight of the Knowledge Guard loss')
    parser.add_argument('--w_indep', default=None, type=float, help='Weight of the feature-independence loss')
    parser.add_argument('--margin', default=None, type=float, help='TAR repulsion margin')
    # Optimizer / training hyperparameters
    parser.add_argument('--lr', default=None, type=float, help='AdamW learning rate')
    parser.add_argument('--weight_decay', default=None, type=float, help='AdamW weight decay')
    parser.add_argument('--label_smoothing', default=None, type=float, help='Cross-entropy label smoothing')
    parser.add_argument('--grad_clip_norm', default=None, type=float, help='Gradient clipping max norm')
    # Model architecture hyperparameters
    parser.add_argument('--embed_dim', default=None, type=int, help='Visual/text embedding dimensionality')
    parser.add_argument('--bottleneck_ratio', default=None, type=int,
                         help='Reduction factor for the Procedure-Scale and Dynamic Text Decoder bottleneck MLPs')
    parser.add_argument('--phase_kernel_sizes', default=None, type=str,
                         help='Comma-separated kernel sizes for the Phase-Level multi-scale causal conv pyramid, e.g. "3,15,31"')
    parser.add_argument('--logit_scale', default=None, type=float, help='Temperature multiplier on cosine-similarity logits')
    return parser.parse_args()


def resolve_hparam(name, cli_value, cfg_model, cfg_training):
    """CLI override > YAML (`model:` then `training:` section) > hardcoded default."""
    if cli_value is not None:
        return cli_value
    if name in cfg_model:
        return cfg_model[name]
    if name in cfg_training:
        return cfg_training[name]
    return DEFAULTS[name]


def main():
    args = get_args()
    with open(args.config) as f:
        cfg = yaml.safe_load(f)

    cfg_model = cfg.get('model', {})
    cfg_training = cfg.get('training', {})

    hp = {}
    for name in DEFAULTS:
        hp[name] = resolve_hparam(name, getattr(args, name, None), cfg_model, cfg_training)
    if isinstance(hp['phase_kernel_sizes'], str):
        hp['phase_kernel_sizes'] = [int(k) for k in hp['phase_kernel_sizes'].split(',')]

    print(f"Resolved hyperparameters: {hp}")

    # 1. Load the frozen VLM backbone
    surgvlp_cfg = Config.fromfile(args.surgvlp_config)['config']
    raw_model, preprocess = surgvlp.load(surgvlp_cfg.model_config, device=DEVICE)

    # 2. Build federated data loaders for this dataset (config-driven)
    train_loaders, test_loaders = build_loaders(cfg, preprocess, hp['frac'])

    # 3. Encode class prompts with the frozen text encoder
    with open(cfg['class_prompt']) as f:
        class_texts = [line.strip() for line in f.readlines()]
    class_tokens = surgvlp.tokenize(class_texts, device=DEVICE)
    with torch.no_grad():
        text_features = raw_model(None, class_tokens, mode='text')['text_emb']
        text_features /= text_features.norm(dim=-1, keepdim=True)

    # 4. Build the Hi-FL model and initialize the text decoder with CLIP prototypes
    global_model = HiFLModel(
        raw_model,
        num_classes=text_features.shape[0],
        embed_dim=hp['embed_dim'],
        bottleneck_ratio=hp['bottleneck_ratio'],
        phase_kernel_sizes=tuple(hp['phase_kernel_sizes']),
        logit_scale=hp['logit_scale'],
    ).to(DEVICE)
    global_model.init_with_clip(text_features)

    # 5. Federated training with TAR centroid exchange
    loss_weights = {k: hp[k] for k in
                     ('w_tar', 'w_kg', 'w_indep', 'margin', 'lr', 'weight_decay', 'label_smoothing', 'grad_clip_norm')}
    best, history = run_federated_training(
        global_model, train_loaders, test_loaders, cfg['centers'], DEVICE,
        loss_weights, num_rounds=hp['rounds']
    )

    print(f"\nBest result for {cfg['dataset_name']}: acc={best['acc']:.4f} f1={best['f1']:.4f} (round {best['round']})")

    out_path = args.out or f"{cfg['dataset_name']}_best_global_model.pt"
    torch.save(global_model.state_dict(), out_path)
    print(f"Saved best global model to {out_path}")


if __name__ == "__main__":
    main()