diff --git a/README.md b/README.md index 2604349..55ac00e 100644 --- a/README.md +++ b/README.md @@ -18,6 +18,7 @@ WaterFlow/ │ └── utils.py # Metrics, plotting, logging utilities ├── scripts/ # Executable scripts │ ├── train.py # Training pipeline +│ ├── train_confidence.py # Train the confidence scorer on cached candidates │ ├── inference.py # Run inference on trained models │ ├── cache_candidates.py # Sample candidate waters for confidence training │ ├── generate_esm_embeddings.py # Precompute ESM embeddings @@ -26,6 +27,7 @@ WaterFlow/ │ ├── test_dataset.py # Dataset and preprocessing tests │ ├── test_distributed.py # DDP helper and cache prebuild tests │ ├── test_confidence.py # Confidence scorer, target and clustering tests +│ ├── test_train_confidence.py # Confidence trainer: loss, freezing, epoch │ ├── test_flow.py # Flow matching tests │ ├── test_encoder.py # Encoder tests │ ├── test_forward.py # End-to-end forward pass tests @@ -320,6 +322,32 @@ uv run torchrun --nproc_per_node=4 -m scripts.train \ --batch_size 4 # per rank -> effective 16 ``` +### Confidence Model Training + +Trains `ConfidenceGVP` to score flow-sampled candidate waters, reusing the flow +run's cache layout and config plus a per-PDB candidate directory: + +```bash +uv run python -m scripts.train_confidence \ + --flow_run_dir \ + --train_list splits/conf_train.txt \ + --val_list splits/conf_valid.txt \ + --candidate_dir \ + --processed_dir \ + --base_pdb_dir \ + --save_dir \ + --run_name \ + --init_from /checkpoints/best.pt --freeze_backbone +``` + +`--init_from` warm-starts the shared backbone from a flow checkpoint; +`--freeze_backbone` then trains only the score head. Validation reports AUC-PR +(for checkpoint selection) and best F1. Multi-GPU works exactly like flow +training — prefix with `torchrun --nproc_per_node=N`, no flag needed: each rank +trains a disjoint shard, the loss is all-reduced, and the (score, label) pairs +are pooled across ranks so AUC-PR/F1 rank the full candidate set. Rank 0 alone +writes checkpoints. + ### Resuming from Checkpoints To resume training from a checkpoint, you can load the model weights and optimizer state: diff --git a/pyproject.toml b/pyproject.toml index 282182b..6f15924 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,6 +13,7 @@ dependencies = [ "biotite", "pymol-open-source-whl>=3.1.0.4", "scipy", + "scikit-learn", "pandas", "numpy", "matplotlib", diff --git a/scripts/inference.py b/scripts/inference.py index 8f18ecb..7d13f0c 100644 --- a/scripts/inference.py +++ b/scripts/inference.py @@ -29,7 +29,7 @@ from src.constants import DEFAULT_EDGE_CUTOFF, NUM_RBF from src.dataset import ProteinWaterDataset -from src.encoder_base import build_encoder +from src.encoder_base import build_encoder, resolve_encoder_config from src.flow import FlowMatcher, FlowWaterGVP from src.utils import ( compute_placement_metrics, @@ -241,25 +241,7 @@ def build_model_from_config(config: dict, device: torch.device) -> nn.Module: Returns: FlowWaterGVP model instance """ - # Use resolved_encoder_config if available (from training), otherwise build from config - resolved = config.get("resolved_encoder_config") - if resolved: - encoder_config = resolved.copy() - else: - encoder_type = config.get("encoder_type", "gvp") - encoder_config = { - "encoder_type": encoder_type, - "hidden_s": config.get("hidden_s") or 256, - "hidden_v": config.get("hidden_v") or 64, - "node_scalar_in": config.get("node_scalar_in") or 16, - "freeze_encoder": config.get("freeze_encoder", False), - "encoder_ckpt": config.get("encoder_ckpt"), - } - - if encoder_type in {"slae", "esm"}: - encoder_config["embedding_key"] = "embedding" - encoder_config["embedding_dim"] = config.get("embedding_dim") - + encoder_config = resolve_encoder_config(config) encoder = build_encoder(encoder_config, device) model = FlowWaterGVP( diff --git a/scripts/train_confidence.py b/scripts/train_confidence.py new file mode 100644 index 0000000..ffd43b7 --- /dev/null +++ b/scripts/train_confidence.py @@ -0,0 +1,824 @@ +#!/usr/bin/env python +""" +Train `ConfidenceGVP` on candidates sampled from a trained flow checkpoint. + +Stage two of the pipeline, riding on the flow dataset/cache layout: + 1. A candidate generator samples waters from a trained flow checkpoint and + writes a thin per-PDB file (`.pt` holding only `candidate_pos`). + 2. This script fits a head with BCE-with-logits on a soft smootherstep target + of each candidate's nearest-GT distance (or a hard 1[d<=cutoff] label), + loading the protein graph, embeddings, and GT waters straight from the + flow caches via `ProteinWaterDataset` + `ConfidenceDataset`. + +`best.pt` is selected on candidate-level AUC-PR at the acceptance label, not on +val loss. Architecture, encoder plumbing, and dataset filters all come from the +flow run's `config.json`, which keeps the two stages in lockstep. +""" + +from __future__ import annotations + +import argparse +import json +import math +from contextlib import nullcontext +from pathlib import Path + +import torch +import torch.nn.functional as F +from loguru import logger +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.optim import AdamW +from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, ReduceLROnPlateau +from torch.utils.data import DataLoader +from torch.utils.data.distributed import DistributedSampler +from torch_geometric.data import Batch +from tqdm import tqdm + +from scripts.inference import _extract_dataset_filter_config +from src.confidence import ConfidenceGVP +from src.confidence_dataset import ConfidenceDataset +from src.constants import NUM_RBF +from src.dataset import ProteinWaterDataset +from src.distributed import ( + all_gather_concat, + all_reduce_means, + ddp_barrier, + ddp_is_active, + is_main_process, + setup_distributed, + teardown_distributed, +) +from src.encoder_base import build_encoder, resolve_encoder_config +from src.utils import auc_pr_and_best_f1, setup_logging_for_tqdm + + +# Everything except the score head. Frozen together, or not at all. +BACKBONE_MODULE_NAMES = ( + "encoder", + "encoder_to_flow", + "protein_scalar_encoder", + "water_scalar_encoder", + "updater", +) + + +def load_config(run_dir: Path) -> dict: + """ + Read a flow run's recorded `config.json`. + + Args: + run_dir: The flow training run directory. + + Returns: + The parsed config. + """ + config_path = run_dir / "config.json" + if not config_path.exists(): + raise FileNotFoundError(f"Config file not found: {config_path}") + with open(config_path) as f: + return json.load(f) + + +def _unwrap(model: ConfidenceGVP | DDP) -> ConfidenceGVP: + """Return the underlying ConfidenceGVP whether or not it is DDP-wrapped.""" + return model.module if isinstance(model, DDP) else model + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser( + description="Train a confidence model from a pre-built candidate cache." + ) + # data / cache paths + p.add_argument( + "--flow_run_dir", + type=str, + required=True, + help="Flow training run directory (provides encoder/filter config).", + ) + p.add_argument("--train_list", type=str, required=True) + p.add_argument("--val_list", type=str, required=True) + p.add_argument( + "--candidate_dir", + type=str, + required=True, + help="Per-PDB candidate directory (.pt with `candidate_pos`), " + "written by scripts/cache_candidates.py.", + ) + p.add_argument( + "--processed_dir", + type=str, + required=True, + help="Cache root shared with flow training (geometry + esm). " + "Protein graph, embeddings, and GT waters load from here.", + ) + p.add_argument( + "--base_pdb_dir", + type=str, + required=True, + help="Base PDB dir, as used by flow training.", + ) + p.add_argument( + "--geometry_cache_name", + type=str, + default=None, + help="Override geometry cache base name (default: flow config).", + ) + p.add_argument( + "--include_mates", + action="store_true", + default=None, + help="Force-include symmetry mates (default: flow config).", + ) + p.add_argument( + "--strict_cache", + action="store_true", + help="Raise when a structure has no candidate file. Default: skip " + "it (the dataset logs how many were dropped).", + ) + p.add_argument( + "--max_candidates", + type=int, + default=None, + help="Cap candidates per structure (fresh random subsample of the " + "pooled cloud each epoch). Bounds per-step memory at no cost " + "to quality, as candidates are scored independently. Unset " + "scores all of them.", + ) + + # run control + p.add_argument( + "--save_dir", + type=str, + required=True, + help="Parent directory for confidence training runs.", + ) + p.add_argument( + "--run_name", + type=str, + required=True, + help="Run identifier. Outputs go to // " + "(checkpoints, config.json, train.log) and it names the wandb run " + "unless --wandb_run_name overrides.", + ) + p.add_argument( + "--init_from", + type=str, + default=None, + help="Checkpoint to warm-start from (e.g. a flow or joint " + "flow+confidence run).", + ) + p.add_argument( + "--freeze_backbone", + action="store_true", + help="Freeze the pretrained encoder/shared backbone and train only " + "the score head.", + ) + + # optimization + p.add_argument("--epochs", type=int, default=50) + p.add_argument("--batch_size", type=int, default=8) + p.add_argument( + "--grad_accum_steps", + type=int, + default=1, + help="Micro-batches per optimizer step. Effective batch = n_gpus * " + "batch_size * grad_accum_steps. Each micro-batch's mean loss is scaled " + "by 1/grad_accum_steps, so unequal candidate counts weight smaller " + "micro-batches up slightly -- the standard DDP accumulation scheme, as " + "in the flow trainer.", + ) + p.add_argument("--num_workers", type=int, default=4) + p.add_argument("--lr", type=float, default=1e-4) + p.add_argument("--weight_decay", type=float, default=1e-5) + p.add_argument("--grad_clip", type=float, default=1.0) + p.add_argument("--warmup_steps", type=int, default=500) + p.add_argument( + "--eta_min_factor", + type=float, + default=0.01, + help="eta_min for cosine = lr * eta_min_factor.", + ) + p.add_argument( + "--scheduler", + type=str, + default="cosine", + choices=["cosine", "plateau"], + help="Main LR scheduler. Both run after the linear warmup.", + ) + p.add_argument( + "--plateau_factor", + type=float, + default=0.5, + help="ReduceLROnPlateau: lr <- lr * factor on plateau.", + ) + p.add_argument( + "--plateau_patience", + type=int, + default=5, + help="ReduceLROnPlateau: epochs of no val improvement.", + ) + p.add_argument("--plateau_min_lr", type=float, default=1e-7) + + # wandb + p.add_argument( + "--wandb_project", + type=str, + default=None, + help="If set, log to this wandb project. Omit to disable wandb.", + ) + p.add_argument("--wandb_entity", type=str, default=None) + p.add_argument( + "--wandb_run_name", + type=str, + default=None, + help="Defaults to --run_name when omitted.", + ) + + # loss shape -- smootherstep target, 0.5-crossing on the acceptance radius + p.add_argument( + "--r_in", + type=float, + default=0.5, + help="smootherstep plateau radius (A): conf=1 for d<=r_in.", + ) + p.add_argument( + "--r_out", + type=float, + default=1.5, + help="smootherstep floor radius (A): conf=0 for d>=r_out. Crossing " + "sits at (r_in+r_out)/2, width sets steepness; the 0.5/1.5 " + "default puts it on the 1A acceptance.", + ) + p.add_argument( + "--accept_radius", + type=float, + default=1.0, + help="Acceptance radius (A). Defines the binary AUC-PR validation " + "label and the cutoff for --hard_label.", + ) + p.add_argument( + "--hard_label", + action="store_true", + help="Train on 1[d<=accept_radius] instead of the soft " + "smootherstep, yielding a calibrated P(within radius).", + ) + # model / system + p.add_argument("--device", type=str, default="cuda") + p.add_argument("--log_level", type=str, default="INFO") + return p.parse_args() + + +def freeze_backbone(model: ConfidenceGVP) -> None: + """Freeze the pretrained backbone, leaving the score head trainable.""" + for name in BACKBONE_MODULE_NAMES: + for param in getattr(model, name).parameters(): + param.requires_grad = False + for param in model.score_head.parameters(): + param.requires_grad = True + + +def _set_model_mode( + model: ConfidenceGVP, *, training: bool, freeze_backbone_enabled: bool +) -> None: + """Set train/eval mode, holding a frozen backbone in eval while the head trains.""" + model.train(training) + if training and freeze_backbone_enabled: + for name in BACKBONE_MODULE_NAMES: + getattr(model, name).eval() + + +def build_confidence_model(config: dict, device: torch.device) -> ConfidenceGVP: + """ + Instantiate `ConfidenceGVP` from the flow run's hyperparameters. + + Mirroring the flow's encoder and hidden dims is what lets the head + warm-start from a checkpoint that shares the same backbone shape. + + Args: + config: The flow run's recorded config. + device: Device to build on. + + Returns: + The model, on `device`. + """ + encoder = build_encoder(resolve_encoder_config(config), device) + return ConfidenceGVP( + encoder=encoder, + hidden_dims=(config.get("hidden_s") or 256, config.get("hidden_v") or 64), + edge_scalar_dim=config.get("edge_scalar_dim") or NUM_RBF, + layers=config.get("flow_layers") or 3, + drop_rate=config.get("drop_rate", 0.1), + n_message_gvps=config.get("n_message_gvps", 2), + n_update_gvps=config.get("n_update_gvps", 2), + cutoff=config.get("cutoff", 8.0), + max_neighbors=config.get("max_neighbors", 256), + knn_fallback_k=config.get("knn_fallback_k", 8), + # Fixed, not from flow config: candidates are scored with no cached PW + # edges and can land where the radius query leaves them isolated, so the + # knn fallback is always needed (see ConfidenceGVP docstring). + dynamic_edge_policy="knn_if_isolated", + ).to(device) + + +def _warm_start_from( + model: ConfidenceGVP, ckpt_path: Path, device: torch.device +) -> None: + """ + Warm-start the shared backbone from a flow (or confidence) checkpoint. + + A flow checkpoint holds the same backbone as `ConfidenceGVP` plus its own + velocity head. Loading non-strict keeps the matching backbone tensors, + ignores the flow-only head (extra keys), and leaves the score head at its + fresh init (missing keys); any shape-mismatched tensor is skipped too. + + Args: + model: Model to load into, modified in place. + ckpt_path: Checkpoint to read. + device: Map location for the load. + """ + state = torch.load(ckpt_path, map_location=device, weights_only=False) + if isinstance(state, dict) and "model_state_dict" in state: + state = state["model_state_dict"] + + target = model.state_dict() + compatible = { + k: v for k, v in state.items() if k in target and target[k].shape == v.shape + } + # No overlap means a wrong or corrupt --init_from; warm-starting would + # silently train from scratch, so fail loud. + if not compatible: + raise ValueError( + f"Warm-start checkpoint {ckpt_path} shares no tensors with " + f"ConfidenceGVP ({len(state)} checkpoint tensors, " + f"{len(target)} model tensors, none matched by name and shape). " + "Check that --init_from points at a matching flow or confidence run." + ) + missing = model.load_state_dict(compatible, strict=False).missing_keys + logger.info( + f"Warm-started from {ckpt_path}: loaded {len(compatible)} tensors, " + f"{len(missing)} left at fresh init." + ) + + +def _build_loader( + args: argparse.Namespace, + pdb_list: str, + shuffle: bool, + config: dict, + *, + distributed: bool = False, + drop_last: bool = False, +) -> tuple[DataLoader, DistributedSampler | None]: + """ + Build a confidence loader over the flow cache layout plus the candidate files. + + The protein graph, embeddings, and GT waters come from the flow caches via + `ProteinWaterDataset`, using the flow run's encoder type and filters so the + graph matches what the flow model saw; candidates come from `candidate_dir`. + + Args: + args: Parsed CLI arguments. + pdb_list: Split file to load. + shuffle: Shuffle the data. Also selects the train-only candidate cap. + config: The flow run's recorded config. + distributed: Shard across ranks with a `DistributedSampler`. + drop_last: Drop the trailing partial batch. Set on the train loader + under DDP so every rank runs the same number of optimizer steps. + + Returns: + (loader, sampler); sampler is None when not distributed. + """ + include_mates = ( + args.include_mates + if args.include_mates is not None + else config.get("include_mates", False) + ) + ds_kwargs = dict( + pdb_list_file=pdb_list, + processed_dir=args.processed_dir, + base_pdb_dir=args.base_pdb_dir, + encoder_type=config.get("encoder_type", "gvp"), + include_mates=include_mates, + # Also picks the cache directory, so it has to track the flow run. + include_ligands=config.get("include_ligands", True), + geometry_cache_name=args.geometry_cache_name + or config.get("geometry_cache_name", "geometry"), + preprocess=True, + **_extract_dataset_filter_config(config), + ) + ds = ConfidenceDataset( + flow_dataset=ProteinWaterDataset(**ds_kwargs), + candidate_dir=args.candidate_dir, + r_in=args.r_in, + r_out=args.r_out, + hard_label=args.hard_label, + accept_radius=args.accept_radius, + # Cap on the train loader only: it bounds backward memory and the draw is + # i.i.d., while val stays a fixed full set so its metric is comparable. + max_candidates=args.max_candidates if shuffle else None, + strict=args.strict_cache, + ) + sampler = ( + DistributedSampler(ds, shuffle=shuffle, drop_last=drop_last) + if distributed + else None + ) + loader = DataLoader( + ds, + batch_size=args.batch_size, + # DataLoader forbids shuffle=True alongside a sampler; the sampler shuffles. + shuffle=shuffle if sampler is None else False, + sampler=sampler, + drop_last=drop_last, + num_workers=args.num_workers, + pin_memory=True, + persistent_workers=args.num_workers > 0, + collate_fn=lambda b: Batch.from_data_list(b), + ) + return loader, sampler + + +def train_one_epoch( + model: ConfidenceGVP | DDP, + loader: DataLoader, + optimizer: AdamW, + warmup_scheduler, + device: torch.device, + args: argparse.Namespace, + step_counter: int, + wandb_run=None, +) -> tuple[float, int]: + """ + Run one training epoch. + + Args: + model: The model, DDP-wrapped or bare. + loader: Train loader. + optimizer: AdamW over the trainable parameters. + warmup_scheduler: LinearLR stepped per optimizer step, or None. + device: Compute device. + args: Parsed CLI arguments. + step_counter: Optimizer steps taken so far, for warmup and wandb. + wandb_run: Active wandb run on rank 0, or None. + + Returns: + (epoch mean loss over candidates, updated step_counter). + """ + _set_model_mode( + _unwrap(model), training=True, freeze_backbone_enabled=args.freeze_backbone + ) + total_loss, total_n = 0.0, 0 + params = [p for group in optimizer.param_groups for p in group["params"]] + accum = max(1, args.grad_accum_steps) + # Bound to the wrapper, since only DDP defers the gradient all-reduce. + no_sync = model.no_sync if isinstance(model, DDP) else None + n_batches = len(loader) # Equal across ranks: DistributedSampler + drop_last. + + optimizer.zero_grad(set_to_none=True) + pbar = tqdm(loader, desc="train", leave=False) + for micro, batch in enumerate(pbar, start=1): + batch = batch.to(device) + target = batch["water"].target_confidence + is_empty = target.numel() == 0 + # Step every `accum` micro-batches and on the last one, so a trailing + # partial window still steps; the stepping backward runs outside no_sync + # and all-reduces the whole accumulated gradient. + is_boundary = (micro % accum == 0) or (micro == n_batches) + sync_ctx = ( + no_sync() if (no_sync is not None and not is_boundary) else nullcontext() + ) + + with sync_ctx: + # Always forward through the DDP model, even when empty: the gradient + # reducer is armed inside DDP.forward, so a rank that skipped it would + # never all-reduce and would hang the others. On empty batches the + # model returns a connected (0,) tensor, so this backward is a no-op. + preds = model(batch, return_logits=True) + loss = ( + preds.sum() + if is_empty + else F.binary_cross_entropy_with_logits(preds, target) + ) + (loss / accum).backward() + + if is_boundary: + torch.nn.utils.clip_grad_norm_(params, args.grad_clip) + optimizer.step() + optimizer.zero_grad(set_to_none=True) + step_counter += 1 + if warmup_scheduler is not None and step_counter <= args.warmup_steps: + warmup_scheduler.step() + + if is_empty: + continue + + total_loss += loss.item() * target.size(0) + total_n += target.size(0) + pbar.set_postfix(loss=f"{loss.item():.4f}") + # wandb_run is rank-0 only (guarded at init); log once per optimizer step. + if wandb_run is not None and is_boundary: + wandb_run.log( + { + "train/step_loss": loss.item(), + "train/lr": optimizer.param_groups[0]["lr"], + "train/step": step_counter, + }, + step=step_counter, + ) + + # Average the epoch loss across ranks once, so every rank logs the same + # number and picks the best checkpoint from the same value. + means, _ = all_reduce_means({"loss": total_loss}, total_n, device) + return means.get("loss", 0.0), step_counter + + +@torch.no_grad() +def validate( + model: ConfidenceGVP | DDP, + loader: DataLoader, + device: torch.device, + args: argparse.Namespace, +) -> dict[str, float]: + """ + Score the validation split. + + Args: + model: The model, DDP-wrapped or bare. + loader: Val loader. + device: Compute device. + args: Parsed CLI arguments. + + Returns: + loss, mae, auc_pr, and best_f1, identical on every rank. + """ + _set_model_mode( + _unwrap(model), training=False, freeze_backbone_enabled=args.freeze_backbone + ) + total_loss, total_n, abs_err = 0.0, 0, 0.0 + score_chunks: list[torch.Tensor] = [] + label_chunks: list[torch.Tensor] = [] + + # Ranks only sync after the loop (one gather), so skipping an empty batch + # here cannot desync them. + for batch in tqdm(loader, desc="val", leave=False): + batch = batch.to(device) + target = batch["water"].target_confidence + if target.numel() == 0: + continue + preds = model(batch, return_logits=True) + loss = F.binary_cross_entropy_with_logits(preds, target) + # MAE is reported in probability space, for interpretability. + probs = torch.sigmoid(preds) + abs_err += (probs - target).abs().sum().item() + total_loss += loss.item() * target.size(0) + total_n += target.size(0) + score_chunks.append(probs.detach().float().cpu()) + label_chunks.append(batch["water"].within_accept_radius.detach().float().cpu()) + + # Reduce sums so the scheduler and best-checkpoint logic see identical + # metrics on every rank. DistributedSampler pads val with at most + # world_size-1 duplicates, a negligible bias on the mean. + means, _ = all_reduce_means({"loss": total_loss, "mae": abs_err}, total_n, device) + # AUC-PR ranks candidates globally, so pool the pairs rather than the metric. + scores = all_gather_concat( + torch.cat(score_chunks) if score_chunks else torch.empty(0) + ) + labels = all_gather_concat( + torch.cat(label_chunks) if label_chunks else torch.empty(0) + ) + auc_pr, best_f1 = auc_pr_and_best_f1(scores, labels) + return { + "loss": means.get("loss", 0.0), + "mae": means.get("mae", 0.0), + "auc_pr": auc_pr, + "best_f1": best_f1, + } + + +def save_ckpt( + path: Path, + model: ConfidenceGVP, + optimizer: AdamW, + epoch: int, + val_metrics: dict, + args: argparse.Namespace, +) -> None: + """Write an unwrapped checkpoint, so it reloads under single-GPU inference.""" + path.parent.mkdir(parents=True, exist_ok=True) + torch.save( + { + "epoch": epoch, + "model_state_dict": model.state_dict(), + "optimizer_state_dict": optimizer.state_dict(), + "val_metrics": val_metrics, + "args": vars(args), + }, + path, + ) + + +def main() -> None: + args = parse_args() + rank, local_rank, world_size = setup_distributed() + main_proc = is_main_process(rank) + distributed = ddp_is_active() + device = ( + torch.device(f"cuda:{local_rank}") + if distributed + else torch.device(args.device if torch.cuda.is_available() else "cpu") + ) + + run_dir = Path(args.save_dir) / args.run_name + # Rank 0 owns the run dir; the others wait so it exists before they touch it. + if main_proc: + (run_dir / "checkpoints").mkdir(parents=True, exist_ok=True) + ddp_barrier() + setup_logging_for_tqdm( + level=args.log_level, + log_file=str(run_dir / "train.log") if main_proc else None, + ) + if distributed: + logger.info( + f"DDP active: rank {rank}/{world_size} (local_rank {local_rank}), " + f"device {device}." + ) + logger.info(f"Run directory: {run_dir}") + + flow_config = load_config(Path(args.flow_run_dir)) + if main_proc: + with open(run_dir / "config.json", "w") as f: + json.dump( + {"flow_config": flow_config, "confidence_args": vars(args)}, f, indent=2 + ) + + # drop_last on train keeps every rank on the same step count (DDP lockstep). + train_loader, train_sampler = _build_loader( + args, + args.train_list, + True, + flow_config, + distributed=distributed, + drop_last=distributed, + ) + val_loader, _ = _build_loader( + args, + args.val_list, + False, + flow_config, + distributed=distributed, + drop_last=False, + ) + logger.info( + f"Train samples: {len(train_loader.dataset)}, val: {len(val_loader.dataset)}" + ) + + model = build_confidence_model(flow_config, device) + if args.init_from is not None: + _warm_start_from(model, Path(args.init_from), device) + if args.freeze_backbone: + freeze_backbone(model) + logger.info(f"Frozen backbone modules: {', '.join(BACKBONE_MODULE_NAMES)}") + + raw_model = model + trainable_params = [p for p in raw_model.parameters() if p.requires_grad] + if not trainable_params: + raise ValueError( + "No trainable parameters remain after applying freeze settings." + ) + logger.info( + f"Model parameters: trainable={sum(p.numel() for p in trainable_params):,} / " + f"total={sum(p.numel() for p in raw_model.parameters()):,}" + ) + + if distributed: + # broadcast_buffers=False is safe (LayerNorm only, no running stats); + # find_unused_parameters covers --freeze_backbone and variant configs. + model = DDP( + model, + device_ids=[local_rank], + broadcast_buffers=False, + find_unused_parameters=True, + ) + + optimizer = AdamW(trainable_params, lr=args.lr, weight_decay=args.weight_decay) + warmup_scheduler = ( + LinearLR( + optimizer, start_factor=1e-8, end_factor=1.0, total_iters=args.warmup_steps + ) + if args.warmup_steps > 0 + else None + ) + if args.scheduler == "cosine": + main_scheduler = CosineAnnealingLR( + optimizer, T_max=args.epochs, eta_min=args.lr * args.eta_min_factor + ) + else: + main_scheduler = ReduceLROnPlateau( + optimizer, + mode="min", + factor=args.plateau_factor, + patience=args.plateau_patience, + min_lr=args.plateau_min_lr, + ) + + # Rank 0 only, so every downstream wandb_run guard is implicitly a rank guard. + wandb_run = None + if args.wandb_project is not None and main_proc: + import wandb + + wandb_run = wandb.init( + project=args.wandb_project, + entity=args.wandb_entity, + name=args.wandb_run_name or args.run_name, + dir=str(run_dir), + config={ + "confidence_args": vars(args), + "flow_config_summary": { + k: flow_config.get(k) + for k in ( + "encoder_type", + "hidden_s", + "hidden_v", + "flow_layers", + "cutoff", + "max_neighbors", + "n_message_gvps", + "n_update_gvps", + ) + }, + }, + ) + logger.info(f"wandb run: {wandb_run.name} ({wandb_run.url})") + + # best.pt tracks candidate-level AUC-PR at the acceptance label -- the ranking + # metric the deliverable cares about -- not val loss. + best_aucpr = float("-inf") + step_counter = 0 + for epoch in range(args.epochs): + # Reshuffle each rank's shard differently per epoch (DDP requirement). + if train_sampler is not None: + train_sampler.set_epoch(epoch) + train_loss, step_counter = train_one_epoch( + model, + train_loader, + optimizer, + warmup_scheduler, + device, + args, + step_counter, + wandb_run, + ) + # Metrics are all-reduced inside, so every rank steps the scheduler on the + # same value and their learning rates stay in sync. + val_metrics = validate(model, val_loader, device, args) + if step_counter > args.warmup_steps: + if isinstance(main_scheduler, ReduceLROnPlateau): + main_scheduler.step(val_metrics["loss"]) + else: + main_scheduler.step() + lr = optimizer.param_groups[0]["lr"] + + logger.info( + f"epoch {epoch:3d} | lr={lr:.2e} | train={train_loss:.4f} | " + f"val={val_metrics['loss']:.4f} | mae={val_metrics['mae']:.4f} | " + f"auc_pr={val_metrics['auc_pr']:.4f} | f1={val_metrics['best_f1']:.4f}" + ) + if wandb_run is not None: + wandb_run.log( + { + "epoch": epoch, + "train/epoch_loss": train_loss, + "val/loss": val_metrics["loss"], + "val/mae": val_metrics["mae"], + "val/auc_pr": val_metrics["auc_pr"], + "val/best_f1": val_metrics["best_f1"], + "lr": lr, + }, + step=step_counter, + ) + + aucpr = val_metrics["auc_pr"] + # auc_pr is nan when an epoch has no positives to rank; skip those. + is_best = not math.isnan(aucpr) and aucpr > best_aucpr + if main_proc: + ckpt_dir = run_dir / "checkpoints" + save_ckpt( + ckpt_dir / "last.pt", raw_model, optimizer, epoch, val_metrics, args + ) + if is_best: + best_aucpr = aucpr + save_ckpt( + ckpt_dir / "best.pt", raw_model, optimizer, epoch, val_metrics, args + ) + logger.info(f" new best (val auc_pr = {best_aucpr:.4f})") + elif is_best: + best_aucpr = aucpr + + if wandb_run is not None: + wandb_run.finish() + teardown_distributed() + logger.info("Training complete.") + + +if __name__ == "__main__": + main() diff --git a/src/confidence.py b/src/confidence.py index 0cf4292..52ed536 100644 --- a/src/confidence.py +++ b/src/confidence.py @@ -387,7 +387,20 @@ def forward( device = data["protein"].pos.device if "water" not in data.node_types or data["water"].num_nodes == 0: - return torch.zeros(0, device=device) + # No candidates. Return an empty (0,) result that keeps a grad path to + # the score head: under DDP the backward must reach the reducer, or + # this rank skips the all-reduce and hangs ranks that had candidates. + # Route through the water head only -- 0 nodes cannot build edges. + in_features = self.water_scalar_encoder[0].in_features + water_x = ( + data["water"].x + if "water" in data.node_types + else torch.zeros(0, in_features, device=device) + ) + s_w = self.water_scalar_encoder(water_x) # (0, s_h) + v_w = torch.zeros(0, self.hidden_dims[1], 3, device=device) + logits = self.score_head((s_w, v_w)).squeeze(-1) # (0,) + return logits if return_logits else torch.sigmoid(logits) s_all, v_all, pp_edge_attr = self.encoder(data) encoder_input = (s_all, v_all) if self.encoder.output_dims[1] > 0 else s_all diff --git a/src/distributed.py b/src/distributed.py index da18491..ba04e09 100644 --- a/src/distributed.py +++ b/src/distributed.py @@ -128,6 +128,28 @@ def run_once_on_main(work: Callable[[], None], key: str) -> dist.Store | None: return store +def all_gather_concat(t: torch.Tensor) -> torch.Tensor: + """ + Pool a variable-length 1D tensor from every rank into one tensor. + + Ranking metrics need a single global ordering, so unlike a loss they cannot + be recovered from per-rank means. Every rank receives the same pooled result + and computes the same number. Shipped via CPU, so ranks may differ in length. + + Args: + t: This rank's 1D contribution. + + Returns: + The concatenation over ranks, in rank order; `t` when not distributed. + """ + if not ddp_is_active(): + return t + # all_gather_object overwrites every slot, so the placeholder is only a shape. + gathered: list[torch.Tensor] = [torch.empty(0)] * dist.get_world_size() + dist.all_gather_object(gathered, t.cpu()) + return torch.cat(gathered, dim=0) + + def all_reduce_means( sums: dict[str, float], count: int, device: torch.device ) -> tuple[dict[str, float], int]: diff --git a/src/encoder_base.py b/src/encoder_base.py index f2b81b3..71772f7 100644 --- a/src/encoder_base.py +++ b/src/encoder_base.py @@ -63,6 +63,39 @@ def get_encoder_class(name: str) -> type[BaseProteinEncoder]: return _ENCODER_REGISTRY[name] +def resolve_encoder_config(config: dict) -> dict: + """ + Derive the build_encoder config from a training run's recorded config. + + Prefers the run's resolved_encoder_config, else rebuilds it from the + top-level hyperparameters. Shared by the flow and confidence model builders + so both stages resolve the encoder identically. + + Args: + config: A training run's recorded config. + + Returns: + A config dict accepted by build_encoder. + """ + resolved = config.get("resolved_encoder_config") + if resolved: + return resolved.copy() + + encoder_type = config.get("encoder_type", "gvp") + encoder_config = { + "encoder_type": encoder_type, + "hidden_s": config.get("hidden_s") or 256, + "hidden_v": config.get("hidden_v") or 64, + "node_scalar_in": config.get("node_scalar_in") or 16, + "freeze_encoder": config.get("freeze_encoder", False), + "encoder_ckpt": config.get("encoder_ckpt"), + } + if encoder_type in {"slae", "esm"}: + encoder_config["embedding_key"] = "embedding" + encoder_config["embedding_dim"] = config.get("embedding_dim") + return encoder_config + + def build_encoder(config: dict, device: torch.device) -> BaseProteinEncoder: """ Build encoder from configuration dict. diff --git a/src/utils.py b/src/utils.py index 8e3b660..060ac8e 100644 --- a/src/utils.py +++ b/src/utils.py @@ -660,3 +660,34 @@ def save_protein_plot( ax.set_title(f"Step {step}") plt.savefig(f"{save_dir}/step_{step}.png") plt.close() + + +def auc_pr_and_best_f1(scores: Tensor, labels: Tensor) -> tuple[float, float]: + """ + Average precision and best F1 from one sorted-score precision-recall sweep. + + Scores a *ranking*, so it cannot be averaged from per-shard sums the way a + loss can -- pool the candidates first (see `all_gather_concat`), then compute. + best_f1 is the max of 2PR/(P+R) over every score threshold. + + Args: + scores: (N,) candidate scores; higher ranks first. + labels: (N,) binary labels aligned with `scores`. + + Returns: + (auc_pr, best_f1); both nan when there are no positives to rank. + """ + if labels.numel() == 0 or labels.sum() == 0: + return float("nan"), float("nan") + + lab = labels[torch.argsort(scores, descending=True)].double() + tp = torch.cumsum(lab, dim=0) + precision = tp / torch.arange(1, lab.numel() + 1, device=lab.device) + recall = tp / lab.sum() + rec_prev = torch.cat([recall.new_zeros(1), recall[:-1]]) + ap = torch.sum((recall - rec_prev) * precision).item() + + best_f1 = ( + (2.0 * precision * recall / (precision + recall).clamp_min(1e-12)).max().item() + ) + return ap, best_f1 diff --git a/tests/test_distributed.py b/tests/test_distributed.py index d72f0c9..007aa96 100644 --- a/tests/test_distributed.py +++ b/tests/test_distributed.py @@ -14,12 +14,14 @@ import pytest import torch +import torch.distributed as dist from torch.utils.data import SequentialSampler import scripts.train as train from src.dataset import get_dataloader, ProteinWaterDataset from src.distributed import ( _ddp_world_size, + all_gather_concat, all_reduce_means, ddp_barrier, ddp_is_active, @@ -185,6 +187,31 @@ def test_all_reduce_means_preserves_key_order(): assert means == {"z": 1.0, "a": 2.0, "m": 3.0} +# ============== all_gather_concat ============== + + +def test_all_gather_concat_is_identity_without_launcher(monkeypatch): + monkeypatch.delenv("WORLD_SIZE", raising=False) + t = torch.tensor([3.0, 1.0, 2.0]) + assert torch.equal(all_gather_concat(t), t) + + +def test_all_gather_concat_pools_uneven_shards_in_rank_order(launcher_env, monkeypatch): + """Ranks score different numbers of candidates, so lengths differ by rank.""" + launcher_env(3) + shards = [torch.tensor([0.0]), torch.tensor([1.0, 2.0]), torch.tensor([3.0])] + + def fake_all_gather_object(out, _obj): + out[:] = shards + + monkeypatch.setattr(dist, "get_world_size", lambda: 3) + monkeypatch.setattr(dist, "all_gather_object", fake_all_gather_object) + + pooled = all_gather_concat(shards[1]) + + assert torch.equal(pooled, torch.tensor([0.0, 1.0, 2.0, 3.0])) + + # ============== Reading model config through the DDP wrapper ============== diff --git a/tests/test_train_confidence.py b/tests/test_train_confidence.py new file mode 100644 index 0000000..c9bf949 --- /dev/null +++ b/tests/test_train_confidence.py @@ -0,0 +1,201 @@ +"""Unit tests for scripts/train_confidence.py -- backbone freezing and the training step.""" + +import argparse + +import pytest +import torch +import torch.nn.functional as F +from torch_geometric.data import Batch, HeteroData + +from scripts.train_confidence import ( + _warm_start_from, + BACKBONE_MODULE_NAMES, + freeze_backbone, + train_one_epoch, +) +from src.confidence import ConfidenceGVP +from src.constants import EDGE_PP, NUM_RBF + + +def _confidence_graph(n_prot=8, n_cand=4): + """One graph shaped like a ConfidenceDataset sample.""" + data = HeteroData() + data["protein"].x = F.one_hot( + torch.randint(0, 16, (n_prot,)), num_classes=16 + ).float() + data["protein"].pos = torch.randn(n_prot, 3) + data["protein"].num_nodes = n_prot + data["water"].x = F.one_hot(torch.full((n_cand,), 2), num_classes=16).float() + data["water"].pos = torch.randn(n_cand, 3) + data["water"].num_nodes = n_cand + data["water"].target_confidence = torch.rand(n_cand) + data["water"].within_accept_radius = (torch.arange(n_cand) % 2).float() + data[EDGE_PP].edge_index = torch.tensor([[0, 1, 2], [1, 2, 3]]) + data[EDGE_PP].edge_unit_vectors = torch.randn(3, 3) + data[EDGE_PP].edge_rbf = torch.randn(3, NUM_RBF) + return data + + +def _train_args(**overrides): + args = argparse.Namespace( + grad_clip=1.0, + warmup_steps=0, + freeze_backbone=True, + grad_accum_steps=1, + ) + return argparse.Namespace(**{**vars(args), **overrides}) + + +@pytest.mark.unit +class TestTrainConfidenceFreezing: + def test_freeze_backbone_leaves_only_score_head_trainable( + self, device, gvp_encoder + ): + model = ConfidenceGVP(encoder=gvp_encoder, hidden_dims=(64, 8), layers=1).to( + device + ) + + freeze_backbone(model) + + trainable = { + name for name, param in model.named_parameters() if param.requires_grad + } + assert trainable + assert all(name.startswith("score_head.") for name in trainable) + + def test_train_one_epoch_frozen_backbone_trains_only_score_head( + self, device, gvp_encoder + ): + batch = Batch.from_data_list([_confidence_graph()]).to(device) + model = ConfidenceGVP(encoder=gvp_encoder, hidden_dims=(64, 8), layers=1).to( + device + ) + freeze_backbone(model) + optimizer = torch.optim.AdamW( + [p for p in model.parameters() if p.requires_grad], lr=1e-4 + ) + + # Grads are zeroed after the step, so compare weights, not .grad. + head_before = [p.detach().clone() for p in model.score_head.parameters()] + backbone_before = [p.detach().clone() for p in model.updater.parameters()] + + train_one_epoch( + model=model, + loader=[batch], + optimizer=optimizer, + warmup_scheduler=None, + device=device, + args=_train_args(), + step_counter=0, + ) + + assert model.score_head.training + for name in BACKBONE_MODULE_NAMES: + assert not getattr(model, name).training + + assert any( + not torch.equal(before, param) + for before, param in zip(head_before, model.score_head.parameters()) + ) + assert all( + torch.equal(before, param) + for before, param in zip(backbone_before, model.updater.parameters()) + ) + + def test_unfrozen_backbone_stays_in_train_mode(self, device, gvp_encoder): + batch = Batch.from_data_list([_confidence_graph()]).to(device) + model = ConfidenceGVP(encoder=gvp_encoder, hidden_dims=(64, 8), layers=1).to( + device + ) + optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4) + backbone_before = [p.detach().clone() for p in model.updater.parameters()] + + train_one_epoch( + model=model, + loader=[batch], + optimizer=optimizer, + warmup_scheduler=None, + device=device, + args=_train_args(freeze_backbone=False), + step_counter=0, + ) + + assert model.updater.training + assert any( + not torch.equal(before, param) + for before, param in zip(backbone_before, model.updater.parameters()) + ) + + +@pytest.mark.unit +class TestTrainConfidenceEmptyBatch: + """An empty candidate set must not crash or desync DDP -- the model returns a + grad-connected (0,) tensor so every rank runs a real forward+backward.""" + + def test_forward_on_empty_is_grad_connected(self, device, gvp_encoder): + model = ConfidenceGVP(encoder=gvp_encoder, hidden_dims=(64, 8), layers=1).to( + device + ) + batch = Batch.from_data_list([_confidence_graph(n_cand=0)]).to(device) + + preds = model(batch, return_logits=True) + + assert preds.shape == (0,) + assert preds.requires_grad and preds.grad_fn is not None + preds.sum().backward() + assert any(p.grad is not None for p in model.score_head.parameters()) + + def test_epoch_steps_through_empty_and_nonempty_batches(self, device, gvp_encoder): + model = ConfidenceGVP(encoder=gvp_encoder, hidden_dims=(64, 8), layers=1).to( + device + ) + optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4) + head_before = [p.detach().clone() for p in model.score_head.parameters()] + loader = [ + Batch.from_data_list([_confidence_graph(n_cand=0)]).to(device), + Batch.from_data_list([_confidence_graph()]).to(device), + ] + + _, step_counter = train_one_epoch( + model=model, + loader=loader, + optimizer=optimizer, + warmup_scheduler=None, + device=device, + args=_train_args(freeze_backbone=False), + step_counter=0, + ) + + assert step_counter == 2 # both batches stepped + assert any( + not torch.equal(before, param) + for before, param in zip(head_before, model.score_head.parameters()) + ) + + +@pytest.mark.unit +class TestWarmStart: + def test_raises_when_no_tensors_match(self, tmp_path, device, gvp_encoder): + model = ConfidenceGVP(encoder=gvp_encoder, hidden_dims=(64, 8), layers=1).to( + device + ) + ckpt = tmp_path / "bogus.pt" + torch.save({"unrelated.weight": torch.zeros(3)}, ckpt) + + with pytest.raises(ValueError, match="shares no tensors"): + _warm_start_from(model, ckpt, device) + + def test_loads_matching_backbone(self, tmp_path, device, gvp_encoder): + model = ConfidenceGVP(encoder=gvp_encoder, hidden_dims=(64, 8), layers=1).to( + device + ) + ckpt = tmp_path / "self.pt" + torch.save({"model_state_dict": model.state_dict()}, ckpt) + + # A fresh model warm-started from the first must match it tensor-for-tensor. + other = ConfidenceGVP(encoder=gvp_encoder, hidden_dims=(64, 8), layers=1).to( + device + ) + _warm_start_from(other, ckpt, device) + for (k, a), b in zip(model.state_dict().items(), other.state_dict().values()): + assert torch.equal(a, b), k diff --git a/tests/test_utils.py b/tests/test_utils.py index 45c1145..bd11014 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -12,6 +12,7 @@ All test cases created with assistance from Claude Code and refined. """ +import math from pathlib import Path import biotite.structure as bts @@ -26,6 +27,7 @@ from src.utils import ( + auc_pr_and_best_f1, compute_edge_features, compute_edge_geometry, compute_placement_metrics, @@ -680,3 +682,79 @@ def test_different_sizes(self, tmp_path): # ) # assert Path(gif_path).exists() + + +@pytest.mark.unit +class TestAucPrAndBestF1: + def test_perfect_ranking_scores_one(self): + ap, best_f1 = auc_pr_and_best_f1( + torch.tensor([0.9, 0.8, 0.2, 0.1]), torch.tensor([1.0, 1.0, 0.0, 0.0]) + ) + + assert ap == pytest.approx(1.0) + assert best_f1 == pytest.approx(1.0) + + def test_only_the_order_matters(self): + """A monotone rescale of the scores must not move the metrics.""" + labels = torch.tensor([1.0, 0.0, 1.0, 0.0]) + base = auc_pr_and_best_f1(torch.tensor([0.9, 0.7, 0.5, 0.3]), labels) + rescaled = auc_pr_and_best_f1(torch.tensor([90.0, 7.0, 0.05, 3e-3]), labels) + + assert base == pytest.approx(rescaled) + + def test_no_positives_is_nan(self): + ap, best_f1 = auc_pr_and_best_f1( + torch.tensor([0.9, 0.1]), torch.tensor([0.0, 0.0]) + ) + + assert math.isnan(ap) and math.isnan(best_f1) + + def test_empty_input_is_nan(self): + ap, best_f1 = auc_pr_and_best_f1(torch.empty(0), torch.empty(0)) + + assert math.isnan(ap) and math.isnan(best_f1) + + def test_matches_hand_computed_curve(self): + # Ranked labels [1, 0, 1]: precision 1.0 at recall 0.5, 2/3 at recall 1.0; + # F1 over thresholds = 2/3, 1/2, 4/5 -> best 4/5. + ap, best_f1 = auc_pr_and_best_f1( + torch.tensor([0.9, 0.6, 0.3]), torch.tensor([1.0, 0.0, 1.0]) + ) + + assert ap == pytest.approx(0.5 * 1.0 + 0.5 * (2 / 3)) + assert best_f1 == pytest.approx(4 / 5) + + def test_worse_ranking_scores_lower(self): + """The selection signal: best.pt is chosen on AUC-PR.""" + scores = torch.tensor([0.9, 0.6, 0.3]) + best, _ = auc_pr_and_best_f1(scores, torch.tensor([1.0, 1.0, 0.0])) + worst, _ = auc_pr_and_best_f1(scores, torch.tensor([0.0, 1.0, 1.0])) + + assert best == pytest.approx(1.0) + assert worst < best + + def test_matches_sklearn(self): + """Pin the torch implementation to sklearn: average_precision_score and + the best F1 over its PR curve. The in-loop metric stays torch-native; + this guards it against drift.""" + from sklearn.metrics import average_precision_score, precision_recall_curve + + torch.manual_seed(0) + for _ in range(20): + n = int(torch.randint(5, 50, (1,)).item()) + scores = torch.rand(n) + labels = (torch.rand(n) < 0.4).float() + if labels.sum() == 0: # our contract returns nan; sklearn is undefined + continue + + ap, best_f1 = auc_pr_and_best_f1(scores, labels) + + s = scores.numpy() + y = labels.numpy() + sk_ap = average_precision_score(y, s) + prec, rec, _ = precision_recall_curve(y, s) + with np.errstate(divide="ignore", invalid="ignore"): + sk_f1 = np.nan_to_num(2 * prec * rec / (prec + rec)).max() + + assert ap == pytest.approx(sk_ap, abs=1e-6) + assert best_f1 == pytest.approx(sk_f1, abs=1e-6) diff --git a/uv.lock b/uv.lock index 0017e0e..7c3829c 100644 --- a/uv.lock +++ b/uv.lock @@ -1946,6 +1946,7 @@ dependencies = [ { name = "pillow" }, { name = "pyg-lib" }, { name = "pymol-open-source-whl" }, + { name = "scikit-learn" }, { name = "scipy" }, { name = "torch" }, { name = "torch-cluster" }, @@ -1977,6 +1978,7 @@ requires-dist = [ { name = "pillow" }, { name = "pyg-lib", specifier = ">=0.6.0", index = "https://data.pyg.org/whl/torch-2.8.0+cu126.html" }, { name = "pymol-open-source-whl", specifier = ">=3.1.0.4" }, + { name = "scikit-learn" }, { name = "scipy" }, { name = "torch", specifier = "==2.8.0", index = "https://download.pytorch.org/whl/cu126" }, { name = "torch-cluster", index = "https://data.pyg.org/whl/torch-2.8.0+cu126.html" },