"""Train the frozen-Omni EasyCom or Ego4D anticipation head. The defaults follow EgoSpeak's EasyCom LSTR optimization config where the architectures are comparable: 50 epochs, Adam, 7e-5 learning rate, 5e-5 weight decay, linear warmup for 10 epochs, cosine decay, dropout 0.2, and MCE loss. """ from __future__ import annotations import argparse import json import math import random from pathlib import Path from typing import Any import numpy as np import torch from torch.optim.lr_scheduler import LambdaLR from torch.utils.data import DataLoader, TensorDataset from tqdm.auto import tqdm from easycom_head import ( CLASS_NAMES, DATASET_NAMES, LAYER_FUSION_MODES, LOSS_NAMES, NUM_HORIZONS, OFFICIAL_TRAIN_SESSIONS, STEP_SECONDS, TRANSITION_OBJECTIVE_NAMES, TurnPredictionHead, average_precision_metrics, egospeak_multihorizon_loss, four_state_multihorizon_loss, head_output_names_for_loss, independent_bce_multihorizon_loss, load_ego4d_manifest_split, load_session_data, marginal_probabilities_from_logits, miniroad_multihorizon_loss, normalize_sample_name, output_classes_for_loss, output_space_for_loss, transition_aware_objective, transition_targets_from_activity, validate_disjoint_sessions, ) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--feature-root", type=Path, required=True) parser.add_argument("--target-root", type=Path, required=True) parser.add_argument("--output", type=Path, required=True) parser.add_argument( "--dataset", choices=DATASET_NAMES, default="easycom", help="Dataset naming/split contract (default: easycom).", ) parser.add_argument( "--manifest", type=Path, help="Required for Ego4D; TSV containing split and uid columns.", ) parser.add_argument( "--sessions", nargs="+", default=None, help=( "Training samples. Defaults to the official EasyCom sessions or " "all Ego4D manifest train UIDs." ), ) parser.add_argument( "--dev-sessions", nargs="*", default=[], help=( "Optional held-out training sessions used for checkpoint selection. " "They are automatically removed from --sessions. Never use 1/2/3." ), ) parser.add_argument( "--sparse-anchor-cache-dir", type=Path, action="append", default=[], help=( "Optional transition-sparse AV cache directory containing " "manifest.json. Repeat once per training session. Only its anchor " "target_index values select examples; states still come from " "--feature-root. Development examples always remain dense." ), ) parser.add_argument("--epochs", type=int, default=50) parser.add_argument("--batch-size", type=int, default=512) parser.add_argument("--learning-rate", type=float, default=7e-5) parser.add_argument("--weight-decay", type=float, default=5e-5) parser.add_argument("--warmup-epochs", type=float, default=10.0) parser.add_argument("--warmup-factor", type=float, default=0.3) parser.add_argument("--dropout", type=float, default=0.2) parser.add_argument( "--learnable-layer-mixture", action="store_true", help=( "Legacy alias for --layer-fusion scalar-mixture." ), ) parser.add_argument( "--layer-fusion", choices=tuple(mode for mode in LAYER_FUSION_MODES if mode != "single"), default=None, help=( "Fusion for stacked [T,L,D] query features. fixed-mean has no " "fusion parameters; scalar-mixture learns L global weights; " "concat-mlp flattens L*D and learns a full projection to D." ), ) parser.add_argument( "--loss", choices=LOSS_NAMES, default="lstr-mce", help=( "lstr-mce/miniroad-normalized use three competing softmax logits; " "independent-bce uses three sigmoid marginals; four-state-ce uses " "silence/wearer-only/others-only/overlap softmax logits." ), ) parser.add_argument( "--transition-objective", choices=TRANSITION_OBJECTIVE_NAMES, default="none", help=( "none keeps the activity-only baseline; r1-supervised adds a " "weighted binary transition head; r2-expected-reward additionally " "maximizes exact expected transition correctness." ), ) parser.add_argument("--transition-loss-weight", type=float, default=0.2) parser.add_argument("--transition-pos-weight", type=float, default=3.0) parser.add_argument("--transition-reward-weight", type=float, default=0.1) parser.add_argument("--device", default="cuda:0") parser.add_argument("--num-workers", type=int, default=0) parser.add_argument("--seed", type=int, default=0) parser.add_argument( "--expected-mode", default="causal-prefix", help="Expected feature mode, or 'any' to disable the check.", ) parser.add_argument("--no-progress", action="store_true") return parser.parse_args() def _validate_args(args: argparse.Namespace) -> None: if args.epochs <= 0: raise ValueError("--epochs must be positive") if args.batch_size <= 0: raise ValueError("--batch-size must be positive") if args.learning_rate <= 0: raise ValueError("--learning-rate must be positive") if args.weight_decay < 0: raise ValueError("--weight-decay cannot be negative") if not 0 <= args.warmup_epochs <= args.epochs: raise ValueError("--warmup-epochs must be in [0, epochs]") if not 0 < args.warmup_factor <= 1: raise ValueError("--warmup-factor must be in (0, 1]") if not 0 <= args.dropout < 1: raise ValueError("--dropout must be in [0, 1)") if args.num_workers < 0: raise ValueError("--num-workers cannot be negative") if args.transition_loss_weight < 0: raise ValueError("--transition-loss-weight cannot be negative") if args.transition_pos_weight <= 0: raise ValueError("--transition-pos-weight must be positive") if args.transition_reward_weight < 0: raise ValueError("--transition-reward-weight cannot be negative") if args.dataset == "ego4d" and args.manifest is None: raise ValueError("--manifest is required when --dataset ego4d") if ( args.learnable_layer_mixture and args.layer_fusion not in {None, "scalar-mixture"} ): raise ValueError( "--learnable-layer-mixture is an alias for " "--layer-fusion scalar-mixture and cannot select another fusion" ) def _resolve_device(name: str) -> torch.device: device = torch.device(name) if device.type == "cuda" and not torch.cuda.is_available(): raise RuntimeError(f"CUDA requested but unavailable: {name}") return device def _seed_everything(seed: int) -> None: random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def _load_examples( feature_root: Path, target_root: Path, sessions: list[str], expected_mode: str | None, dataset: str, selected_indices: dict[str, set[int]] | None = None, ) -> tuple[ torch.Tensor, torch.Tensor, torch.Tensor, int, int, list[dict[str, Any]], ]: all_states: list[torch.Tensor] = [] all_targets: list[torch.Tensor] = [] all_transition_targets: list[torch.Tensor] = [] metadata: list[dict[str, Any]] = [] hidden_size: int | None = None num_input_layers: int | None = None for name in sessions: data = load_session_data( feature_root, target_root, name, expected_mode=expected_mode, dataset=dataset, ) states, targets, future_indices = data.future_examples(NUM_HORIZONS) current_indices = future_indices[:, 0] - 1 activity_sequence = torch.cat( (data.targets[current_indices, None, :], targets), dim=1, ) transition_targets = transition_targets_from_activity(activity_sequence) dense_example_count = len(states) if selected_indices is not None: if name not in selected_indices: raise ValueError(f"No sparse anchor manifest was supplied for {name}") wanted = selected_indices[name] keep = torch.tensor( [int(index) in wanted for index in current_indices.tolist()], dtype=torch.bool, ) states = states[keep] targets = targets[keep] transition_targets = transition_targets[keep] if not len(states): raise ValueError(f"Sparse anchor selection kept no examples for {name}") if hidden_size is None: hidden_size = data.hidden_size elif data.hidden_size != hidden_size: raise ValueError( f"Feature dimensions differ: {name} has {data.hidden_size}, " f"expected {hidden_size}" ) if num_input_layers is None: num_input_layers = data.num_input_layers elif data.num_input_layers != num_input_layers: raise ValueError( f"Feature layer counts differ: {name} has " f"{data.num_input_layers}, expected {num_input_layers}" ) all_states.append(states) all_targets.append(targets) all_transition_targets.append(transition_targets) metadata.append( { "session": name, "num_states": len(data.states), "num_training_examples": len(states), "num_dense_training_examples": dense_example_count, "num_positive_transition_targets": int( transition_targets.sum().item() ), "transition_target_count": int(transition_targets.numel()), "transition_positive_rate": float(transition_targets.mean()), "supervision_policy": ( "transition-sparse-manifest" if selected_indices is not None else "dense-5hz" ), **data.feature_metadata, } ) print( f"{name}: states={len(data.states)}, " f"valid_10_horizon_examples={dense_example_count}, " f"selected_examples={len(states)}" ) if hidden_size is None or num_input_layers is None: raise ValueError("No sessions selected") return ( torch.cat(all_states, dim=0), torch.cat(all_targets, dim=0), torch.cat(all_transition_targets, dim=0), hidden_size, num_input_layers, metadata, ) def _load_sparse_anchor_indices( cache_dirs: list[Path], ) -> dict[str, set[int]] | None: if not cache_dirs: return None result: dict[str, set[int]] = {} for cache_dir in cache_dirs: manifest_path = cache_dir.expanduser().resolve() / "manifest.json" if not manifest_path.is_file(): raise FileNotFoundError(manifest_path) manifest = json.loads(manifest_path.read_text()) session = normalize_sample_name(str(manifest.get("session", "")), "easycom") policy = manifest.get("anchor_selection", {}).get("policy") if policy != "transition-sparse": raise ValueError( f"Expected transition-sparse manifest, got {policy!r}: " f"{manifest_path}" ) if session in result: raise ValueError(f"Duplicate sparse anchor manifest for {session}") indices = { int(anchor["target_index"]) for anchor in manifest.get("anchors", []) if bool(anchor.get("has_future_target", True)) } if not indices: raise ValueError(f"No valid sparse anchors in {manifest_path}") result[session] = indices print(f"sparse anchors {session}: {len(indices)} from {manifest_path}") return result def _scheduler( optimizer: torch.optim.Optimizer, *, steps_per_epoch: int, epochs: int, warmup_epochs: float, warmup_factor: float, ) -> LambdaLR: total_steps = max(1, steps_per_epoch * epochs) warmup_steps = steps_per_epoch * warmup_epochs def multiplier(step: int) -> float: cosine = 0.5 * (1.0 + math.cos(math.pi * step / total_steps)) if warmup_steps > 0 and step < warmup_steps: alpha = step / warmup_steps warmup = warmup_factor * (1.0 - alpha) + alpha else: warmup = 1.0 return warmup * cosine return LambdaLR(optimizer, multiplier) @torch.inference_mode() def _evaluate_valid_horizons( head: TurnPredictionHead, states: torch.Tensor, targets: torch.Tensor, device: torch.device, batch_size: int, loss_name: str, ) -> dict[str, Any]: head.eval() chunks: list[torch.Tensor] = [] for start in range(0, len(states), batch_size): x = states[start : start + batch_size].to( device=device, dtype=torch.float32, non_blocking=True, ) chunks.append( marginal_probabilities_from_logits(head(x), loss_name).cpu() ) probabilities = torch.cat(chunks, dim=0).numpy() target_array = targets.numpy() results = [] for horizon in range(NUM_HORIZONS): result = average_precision_metrics( probabilities[:, horizon], target_array[:, horizon], ) result["horizon"] = horizon + 1 result["seconds"] = (horizon + 1) * STEP_SECONDS results.append(result) return { "horizons": results, "mean_map": float(np.mean([item["map"] for item in results])), } def _checkpoint( head: TurnPredictionHead, *, epoch: int, best_epoch: int, best_dev_map: float | None, train_loss: float, args: argparse.Namespace, train_sessions: list[str], dev_sessions: list[str], hidden_size: int, num_input_layers: int, feature_metadata: list[dict[str, Any]], history: list[dict[str, Any]], ) -> dict[str, Any]: return { "format_version": 3, "head": head.state_dict(), "head_config": { "hidden_size": hidden_size, "num_horizons": NUM_HORIZONS, "num_classes": head.num_classes, "dropout": args.dropout, "num_input_layers": num_input_layers, "layer_fusion": args.layer_fusion, "learnable_layer_mixture": args.learnable_layer_mixture, "transition_auxiliary": head.transition_auxiliary, }, "alignment": { "state_label_offset": 0, "future_offsets": list(range(1, NUM_HORIZONS + 1)), "future_seconds": [ (index + 1) * STEP_SECONDS for index in range(NUM_HORIZONS) ], }, "classes": list(CLASS_NAMES), "head_output_names": list(head_output_names_for_loss(args.loss)), "output_space": output_space_for_loss(args.loss), "dataset": args.dataset, "manifest": ( str(args.manifest.expanduser().resolve()) if args.manifest is not None else None ), "epoch": epoch, "best_epoch": best_epoch, "best_dev_map": best_dev_map, "train_loss": train_loss, "train_sessions": train_sessions, "dev_sessions": dev_sessions, "feature_metadata": feature_metadata, "training_config": { "epochs": args.epochs, "batch_size": args.batch_size, "learning_rate": args.learning_rate, "weight_decay": args.weight_decay, "warmup_epochs": args.warmup_epochs, "warmup_factor": args.warmup_factor, "loss": args.loss, "transition_objective": args.transition_objective, "transition_loss_weight": args.transition_loss_weight, "transition_pos_weight": args.transition_pos_weight, "transition_reward_weight": args.transition_reward_weight, "optimizer": "Adam", "scheduler": "linear_warmup_times_cosine", "seed": args.seed, "expected_mode": args.expected_mode, "layer_fusion": args.layer_fusion, "learnable_layer_mixture": args.layer_fusion == "scalar-mixture", "supervision_policy": ( "transition-sparse-manifest" if args.sparse_anchor_cache_dir else "dense-5hz" ), "sparse_anchor_cache_dirs": [ str(path.expanduser().resolve()) for path in args.sparse_anchor_cache_dir ], }, "history": history, } def main() -> None: args = parse_args() _validate_args(args) if args.layer_fusion is None: args.layer_fusion = ( "scalar-mixture" if args.learnable_layer_mixture else "single" ) args.learnable_layer_mixture = args.layer_fusion == "scalar-mixture" _seed_everything(args.seed) device = _resolve_device(args.device) expected_mode = None if args.expected_mode == "any" else args.expected_mode sparse_anchor_indices = _load_sparse_anchor_indices( args.sparse_anchor_cache_dir ) if sparse_anchor_indices is not None and args.dataset != "easycom": raise ValueError("--sparse-anchor-cache-dir currently supports EasyCom only") if args.dataset == "ego4d": assert args.manifest is not None official_train = load_ego4d_manifest_split(args.manifest, "train") official_test = set(load_ego4d_manifest_split(args.manifest, "test")) requested_values = args.sessions or official_train else: official_train = list(OFFICIAL_TRAIN_SESSIONS) official_test = {"Session_1", "Session_2", "Session_3"} requested_values = args.sessions or official_train requested_train = [ normalize_sample_name(item, args.dataset) for item in requested_values ] dev_sessions = [ normalize_sample_name(item, args.dataset) for item in args.dev_sessions ] if len(requested_train) != len(set(requested_train)): raise ValueError("Duplicate training samples were supplied") if len(dev_sessions) != len(set(dev_sessions)): raise ValueError("Duplicate development samples were supplied") if args.dataset == "ego4d": allowed_train = set(official_train) unknown = (set(requested_train) | set(dev_sessions)) - allowed_train if unknown: raise ValueError( "Ego4D training/dev samples must belong to the manifest train " f"split; invalid: {sorted(unknown)}" ) train_sessions = [item for item in requested_train if item not in dev_sessions] validate_disjoint_sessions(train_sessions, dev_sessions) if not train_sessions: raise ValueError("No training sessions remain after removing dev sessions") if sparse_anchor_indices is not None: missing = set(train_sessions) - set(sparse_anchor_indices) extra = set(sparse_anchor_indices) - set(train_sessions) if missing or extra: raise ValueError( "Sparse anchor manifests must match training sessions exactly; " f"missing={sorted(missing)}, extra={sorted(extra)}" ) leaked = official_test & (set(train_sessions) | set(dev_sessions)) if leaked: raise ValueError( f"Official test sessions cannot be used for training/selection: " f"{sorted(leaked)}" ) print(f"device: {device}") print(f"train sessions: {', '.join(train_sessions)}") print(f"dev sessions: {', '.join(dev_sessions) if dev_sessions else 'none'}") ( states, targets, transition_targets, hidden_size, num_input_layers, feature_metadata, ) = _load_examples( args.feature_root, args.target_root, train_sessions, expected_mode, args.dataset, sparse_anchor_indices, ) print( f"training examples: {len(states)}, feature_dim={hidden_size}, " f"input_layers={num_input_layers}" ) print( "transition targets: " f"{int(transition_targets.sum().item())}/{transition_targets.numel()} " f"positive ({float(transition_targets.mean()) * 100:.3f}%)" ) if num_input_layers > 1 and args.layer_fusion == "single": raise ValueError( f"Features contain {num_input_layers} stacked layers; pass " "--layer-fusion fixed-mean, scalar-mixture, or concat-mlp" ) if num_input_layers == 1 and args.layer_fusion != "single": raise ValueError( f"--layer-fusion {args.layer_fusion} requires stacked [T,L,D] " "features" ) dev_states: torch.Tensor | None = None dev_targets: torch.Tensor | None = None if dev_sessions: ( dev_states, dev_targets, _, dev_hidden_size, dev_num_input_layers, _, ) = _load_examples( args.feature_root, args.target_root, dev_sessions, expected_mode, args.dataset, None, ) if dev_hidden_size != hidden_size: raise ValueError( f"Dev feature dimension {dev_hidden_size} differs from " f"training dimension {hidden_size}" ) if dev_num_input_layers != num_input_layers: raise ValueError( f"Dev feature layer count {dev_num_input_layers} differs from " f"training count {num_input_layers}" ) print(f"development examples: {len(dev_states)}") dataset = TensorDataset(states, targets, transition_targets) generator = torch.Generator().manual_seed(args.seed) loader = DataLoader( dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=device.type == "cuda", persistent_workers=args.num_workers > 0, generator=generator, ) num_output_classes = output_classes_for_loss(args.loss) head = TurnPredictionHead( hidden_size=hidden_size, num_horizons=NUM_HORIZONS, num_classes=num_output_classes, dropout=args.dropout, num_input_layers=num_input_layers, learnable_layer_mixture=args.learnable_layer_mixture, layer_fusion=args.layer_fusion, transition_auxiliary=args.transition_objective != "none", ).to(device) trainable_parameters = sum( parameter.numel() for parameter in head.parameters() if parameter.requires_grad ) fusion_parameters = ( sum(parameter.numel() for parameter in head.layer_fusion_net.parameters()) if head.layer_fusion_net is not None else (head.layer_logits.numel() if head.layer_logits is not None else 0) ) print( f"loss={args.loss}, output_space={output_space_for_loss(args.loss)}, " f"transition_objective={args.transition_objective}, " f"activity_outputs={num_output_classes}, " "transition_outputs=" f"{NUM_HORIZONS if args.transition_objective != 'none' else 0}, " f"layer_fusion={args.layer_fusion}, " f"fusion_parameters={fusion_parameters:,}, " f"head_trainable_parameters={trainable_parameters:,}" ) if head.layer_logits is None: optimizer_parameters: Any = head.parameters() else: optimizer_parameters = [ { "params": [ parameter for name, parameter in head.named_parameters() if name != "layer_logits" ], "weight_decay": args.weight_decay, }, { "params": [head.layer_logits], # Weight decay on softmax logits artificially pulls the layer # distribution back toward uniform. "weight_decay": 0.0, }, ] optimizer = torch.optim.Adam( optimizer_parameters, lr=args.learning_rate, weight_decay=args.weight_decay, ) scheduler = _scheduler( optimizer, steps_per_epoch=len(loader), epochs=args.epochs, warmup_epochs=args.warmup_epochs, warmup_factor=args.warmup_factor, ) loss_function = { "lstr-mce": egospeak_multihorizon_loss, "miniroad-normalized": miniroad_multihorizon_loss, "independent-bce": independent_bce_multihorizon_loss, "four-state-ce": four_state_multihorizon_loss, }[args.loss] args.output.parent.mkdir(parents=True, exist_ok=True) history: list[dict[str, Any]] = [] best_dev_map = -math.inf if dev_sessions else None best_epoch = args.epochs if not dev_sessions else 0 for epoch in range(1, args.epochs + 1): head.train() loss_sum = 0.0 activity_loss_sum = 0.0 transition_loss_sum = 0.0 transition_reward_sum = 0.0 example_count = 0 progress = tqdm( loader, desc=f"head epoch {epoch}/{args.epochs}", unit="batch", disable=args.no_progress, ) for state_batch, target_batch, transition_target_batch in progress: state_batch = state_batch.to( device=device, dtype=torch.float32, non_blocking=True, ) target_batch = target_batch.to( device=device, dtype=torch.float32, non_blocking=True, ) transition_target_batch = transition_target_batch.to( device=device, dtype=torch.float32, non_blocking=True, ) optimizer.zero_grad(set_to_none=True) if args.transition_objective == "none": logits = head(state_batch) activity_loss = loss_function(logits, target_batch) loss = activity_loss transition_loss = activity_loss.detach().new_zeros(()) expected_reward = activity_loss.detach().new_zeros(()) else: logits, transition_logits = head.forward_with_transition(state_batch) activity_loss = loss_function(logits, target_batch) components = transition_aware_objective( activity_loss, transition_logits, transition_target_batch, objective=args.transition_objective, transition_weight=args.transition_loss_weight, pos_weight=args.transition_pos_weight, reward_weight=args.transition_reward_weight, ) loss = components["total"] transition_loss = components["transition"] expected_reward = components["expected_reward"] loss.backward() optimizer.step() scheduler.step() batch_count = len(state_batch) loss_sum += float(loss.detach()) * batch_count activity_loss_sum += float(activity_loss.detach()) * batch_count transition_loss_sum += float(transition_loss.detach()) * batch_count transition_reward_sum += float(expected_reward.detach()) * batch_count example_count += batch_count postfix = { "loss": f"{float(loss.detach()):.5f}", "lr": f"{scheduler.get_last_lr()[0]:.3g}", } if args.transition_objective != "none": postfix["transition"] = f"{float(transition_loss.detach()):.5f}" postfix["reward"] = f"{float(expected_reward.detach()):.5f}" progress.set_postfix(postfix) train_loss = loss_sum / example_count record: dict[str, Any] = { "epoch": epoch, "train_loss": train_loss, "train_activity_loss": activity_loss_sum / example_count, "learning_rate": scheduler.get_last_lr()[0], } if args.transition_objective != "none": record["train_transition_loss"] = ( transition_loss_sum / example_count ) record["train_expected_transition_reward"] = ( transition_reward_sum / example_count ) layer_weights = head.layer_mixture_weights() if layer_weights is not None: record["layer_mixture_weights_oldest_to_newest"] = ( layer_weights.detach().float().cpu().tolist() ) weight_text = ( "" if layer_weights is None else ", layer_weights=[" + ",".join( f"{value:.4f}" for value in record["layer_mixture_weights_oldest_to_newest"] ) + "]" ) transition_text = ( "" if args.transition_objective == "none" else ( f", activity_loss={record['train_activity_loss']:.6f}, " f"transition_loss={record['train_transition_loss']:.6f}, " "expected_transition_reward=" f"{record['train_expected_transition_reward']:.6f}" ) ) if dev_sessions: assert dev_states is not None and dev_targets is not None dev_metrics = _evaluate_valid_horizons( head, dev_states, dev_targets, device, args.batch_size, args.loss, ) dev_map = float(dev_metrics["mean_map"]) record["dev_valid_mean_map"] = dev_map print( f"epoch {epoch:02d}: train_loss={train_loss:.6f}, " f"dev_valid_mean_mAP={dev_map * 100:.3f}" f"{transition_text}{weight_text}" ) if dev_map > float(best_dev_map): best_dev_map = dev_map best_epoch = epoch torch.save( _checkpoint( head, epoch=epoch, best_epoch=best_epoch, best_dev_map=best_dev_map, train_loss=train_loss, args=args, train_sessions=train_sessions, dev_sessions=dev_sessions, hidden_size=hidden_size, num_input_layers=num_input_layers, feature_metadata=feature_metadata, history=history + [record], ), args.output, ) print(f"saved new best checkpoint: {args.output}") else: print( f"epoch {epoch:02d}: train_loss={train_loss:.6f}" f"{transition_text}{weight_text}" ) history.append(record) if not dev_sessions: torch.save( _checkpoint( head, epoch=args.epochs, best_epoch=best_epoch, best_dev_map=None, train_loss=history[-1]["train_loss"], args=args, train_sessions=train_sessions, dev_sessions=dev_sessions, hidden_size=hidden_size, num_input_layers=num_input_layers, feature_metadata=feature_metadata, history=history, ), args.output, ) print(f"checkpoint: {args.output}") if dev_sessions: print( f"best epoch: {best_epoch}, " f"dev valid-only mean mAP: {float(best_dev_map) * 100:.3f}" ) summary_path = args.output.with_suffix(args.output.suffix + ".json") payload = torch.load(args.output, map_location="cpu", weights_only=False) summary = {key: value for key, value in payload.items() if key != "head"} summary_path.write_text(json.dumps(summary, indent=2) + "\n") print(f"training summary: {summary_path}") if payload["head_config"].get("layer_fusion") == "scalar-mixture": saved_logits = payload["head"]["layer_logits"].float() saved_weights = saved_logits.softmax(dim=0).tolist() print( "saved layer mixture weights (oldest->newest): " + ", ".join(f"{value:.6f}" for value in saved_weights) ) if __name__ == "__main__": main()