From f45ea9cbc62354cd434b1d87b0f140c2f6b78556 Mon Sep 17 00:00:00 2001 From: Xin Wang Date: Wed, 26 Aug 2026 13:06:18 +0800 Subject: [PATCH] Add higher-capacity last-four-layer fusion --- README.md | 68 +++++++++++ easycom_head.py | 64 ++++++++-- infer_turn_head.py | 1 + .../README.md | 111 ++++++++++++++++++ .../metrics.csv | 12 ++ .../raw_console.txt | 33 ++++++ .../summary.json | 66 +++++++++++ train_omni_turn_lora.py | 3 + train_turn_head.py | 56 +++++++-- 9 files changed, 397 insertions(+), 17 deletions(-) create mode 100644 results/easycom_frozen_3b_8s_last4_learned_mix_2026-08-26/README.md create mode 100644 results/easycom_frozen_3b_8s_last4_learned_mix_2026-08-26/metrics.csv create mode 100644 results/easycom_frozen_3b_8s_last4_learned_mix_2026-08-26/raw_console.txt create mode 100644 results/easycom_frozen_3b_8s_last4_learned_mix_2026-08-26/summary.json diff --git a/README.md b/README.md index fc5a13d..d4d0ffc 100644 --- a/README.md +++ b/README.md @@ -79,6 +79,74 @@ The raw 8-second console output, machine-readable comparison, and detailed interpretation are archived in [`results/easycom_open_stream_8s_vs_20s_2026-08-23/`](results/easycom_open_stream_8s_vs_20s_2026-08-23/README.md). +### Last-4 query-layer learned mixture + +A frozen Qwen2.5-Omni-3B experiment now reads the query states from the final +four Thinker layers and learns four global softmax mixture weights before the +same `10 x 3` anticipation head. With an 8-second open-stream context, 2 FPS +video, 5 Hz outputs, and 25 head-training epochs, it reaches **73.45085** mAP +at 0.2 seconds and **60.53423** mean mAP on official test Sessions 3, 2, and 1. +Relative to the recorded single-final-layer 8-second run, this is **+1.45317** +at 0.2 seconds and **+2.01486** mean mAP, with positive gains at every horizon. + +The final oldest-to-newest layer weights are +`[0.246384, 0.248327, 0.241774, 0.263515]`. Because they remain almost uniform, +the result currently demonstrates the value of multi-layer query readout, not +yet the value of learning the mixture. A fixed uniform last-4 mean is the next +required control. The complete table, machine-readable metrics, provenance, +and naming caveat are archived in +[`results/easycom_frozen_3b_8s_last4_learned_mix_2026-08-26/`](results/easycom_frozen_3b_8s_last4_learned_mix_2026-08-26/README.md). + +#### Higher-capacity last-4 concat-MLP ablation + +`train_turn_head.py` also supports a higher-capacity fusion control. It +flattens each `[4, 2048]` last-layer stack to `8192` dimensions, applies +`LayerNorm -> Linear(8192, 2048) -> GELU`, and then uses the unchanged +anticipation MLP. This adds about 16.8 million fusion parameters, compared +with four logits in the scalar-mixture experiment. Keep the seed and training +recipe fixed for a clean capacity ablation: + +```bash +CUDA_VISIBLE_DEVICES=0 \ +OMP_NUM_THREADS=4 \ +MKL_NUM_THREADS=4 \ +OPENBLAS_NUM_THREADS=4 \ +python train_turn_head.py \ + --feature-root /data1/yfzx/wangx/Datasets/easycom/omni_features/open_stream_2fps_8s_last4_stack \ + --target-root /data1/yfzx/wangx/Datasets/easycom/easycom_target_frame \ + --output checkpoints/easycom_frozen_3b_8s_last4_concat_mlp_25ep.pth \ + --sessions 4 5 6 7 8 9 10 11 12 \ + --epochs 25 \ + --batch-size 512 \ + --learning-rate 7e-5 \ + --weight-decay 5e-5 \ + --warmup-epochs 10 \ + --warmup-factor 0.3 \ + --dropout 0.2 \ + --loss lstr-mce \ + --layer-fusion concat-mlp \ + --expected-mode open-stream-replay \ + --seed 0 \ + --device cuda:0 +``` + +Evaluate with the same dense stacked test features: + +```bash +CUDA_VISIBLE_DEVICES=0 python infer_turn_head.py \ + --feature-root /data1/yfzx/wangx/Datasets/easycom/omni_features/open_stream_2fps_8s_last4_stack \ + --target-root /data1/yfzx/wangx/Datasets/easycom/easycom_target_frame \ + --checkpoint checkpoints/easycom_frozen_3b_8s_last4_concat_mlp_25ep.pth \ + --output-dir results/easycom_frozen_3b_8s_last4_concat_mlp_25ep \ + --expected-mode open-stream-replay \ + --batch-size 1024 \ + --device cuda:0 +``` + +Treat this as a capacity ablation. If it improves, compare it against a +parameter-matched single-layer MLP and use stronger regularization or a clean +development session before attributing the gain to cross-layer fusion. + ## Current Ego4D context-length results The following single-seed snapshot was recorded on 2026-08-22. Both variants diff --git a/easycom_head.py b/easycom_head.py index 0802c7c..0dc842b 100644 --- a/easycom_head.py +++ b/easycom_head.py @@ -34,6 +34,12 @@ STEP_SECONDS = 0.2 NUM_HORIZONS = 10 NUM_CLASSES = 3 DATASET_NAMES = ("easycom", "ego4d") +LAYER_FUSION_MODES = ( + "single", + "fixed-mean", + "scalar-mixture", + "concat-mlp", +) DatasetName = Literal["easycom", "ego4d"] @@ -48,30 +54,55 @@ class TurnPredictionHead(nn.Module): dropout: float = 0.1, num_input_layers: int = 1, learnable_layer_mixture: bool = False, + layer_fusion: str | None = None, ) -> None: super().__init__() if num_input_layers <= 0: raise ValueError("num_input_layers must be positive") - if learnable_layer_mixture and num_input_layers < 2: - raise ValueError( - "learnable_layer_mixture requires at least two input layers" + if layer_fusion is None: + layer_fusion = ( + "scalar-mixture" if learnable_layer_mixture else "single" ) - if not learnable_layer_mixture and num_input_layers != 1: + if layer_fusion not in LAYER_FUSION_MODES: raise ValueError( - "Multiple input layers require learnable_layer_mixture=True" + f"Unknown layer_fusion={layer_fusion!r}; expected one of " + f"{LAYER_FUSION_MODES}" + ) + if learnable_layer_mixture and layer_fusion != "scalar-mixture": + raise ValueError( + "learnable_layer_mixture=True is only compatible with " + "layer_fusion='scalar-mixture'" + ) + if layer_fusion != "single" and num_input_layers < 2: + raise ValueError( + f"layer_fusion={layer_fusion!r} requires at least two input layers" + ) + if layer_fusion == "single" and num_input_layers != 1: + raise ValueError( + "Multiple input layers require a non-single layer_fusion" ) self.hidden_size = hidden_size self.num_horizons = num_horizons self.num_classes = num_classes self.dropout = dropout self.num_input_layers = num_input_layers - self.learnable_layer_mixture = learnable_layer_mixture - if learnable_layer_mixture: + self.layer_fusion = layer_fusion + self.learnable_layer_mixture = layer_fusion == "scalar-mixture" + if self.learnable_layer_mixture: # Zero initialization exactly reproduces a uniform layer mean at # step zero while allowing training to select useful query layers. self.layer_logits = nn.Parameter(torch.zeros(num_input_layers)) else: self.register_parameter("layer_logits", None) + if layer_fusion == "concat-mlp": + fusion_input_size = num_input_layers * hidden_size + self.layer_fusion_net: nn.Module | None = nn.Sequential( + nn.LayerNorm(fusion_input_size), + nn.Linear(fusion_input_size, hidden_size), + nn.GELU(), + ) + else: + self.layer_fusion_net = None self.net = nn.Sequential( nn.LayerNorm(hidden_size), nn.Linear(hidden_size, hidden_size // 2), @@ -85,13 +116,16 @@ class TurnPredictionHead(nn.Module): raise ValueError( f"Expected hidden size {self.hidden_size}, got {states.shape[-1]}" ) - if self.learnable_layer_mixture: + if self.layer_fusion != "single": if states.ndim < 3 or states.shape[-2] != self.num_input_layers: raise ValueError( - "Layer-mixture input must end in " + "Multi-layer input must end in " f"[{self.num_input_layers},{self.hidden_size}], got " f"{tuple(states.shape)}" ) + if self.layer_fusion == "fixed-mean": + states = states.mean(dim=-2) + elif self.layer_fusion == "scalar-mixture": assert self.layer_logits is not None weights = self.layer_logits.softmax(dim=0).to(states.dtype) view_shape = (1,) * (states.ndim - 2) + ( @@ -99,6 +133,10 @@ class TurnPredictionHead(nn.Module): 1, ) states = (states * weights.view(view_shape)).sum(dim=-2) + elif self.layer_fusion == "concat-mlp": + assert self.layer_fusion_net is not None + states = states.flatten(start_dim=-2) + states = self.layer_fusion_net(states) logits = self.net(states) return logits.view( *states.shape[:-1], self.num_horizons, self.num_classes @@ -406,6 +444,14 @@ def checkpoint_head_config(payload: dict[str, Any]) -> dict[str, Any]: config.setdefault("dropout", 0.1) config.setdefault("num_input_layers", 1) config.setdefault("learnable_layer_mixture", False) + config.setdefault( + "layer_fusion", + ( + "scalar-mixture" + if config["learnable_layer_mixture"] + else "single" + ), + ) return config diff --git a/infer_turn_head.py b/infer_turn_head.py index 582c9ee..8f4386d 100644 --- a/infer_turn_head.py +++ b/infer_turn_head.py @@ -328,6 +328,7 @@ def main() -> None: "checkpoint_epoch": checkpoint_payload.get("epoch"), "checkpoint_train_sessions": checkpoint_payload.get("train_sessions"), "checkpoint_dev_sessions": checkpoint_payload.get("dev_sessions"), + "layer_fusion": head.layer_fusion, "layer_mixture_weights_oldest_to_newest": ( head.layer_mixture_weights().detach().float().cpu().tolist() if head.layer_mixture_weights() is not None diff --git a/results/easycom_frozen_3b_8s_last4_learned_mix_2026-08-26/README.md b/results/easycom_frozen_3b_8s_last4_learned_mix_2026-08-26/README.md new file mode 100644 index 0000000..5148984 --- /dev/null +++ b/results/easycom_frozen_3b_8s_last4_learned_mix_2026-08-26/README.md @@ -0,0 +1,111 @@ +# EasyCom frozen Qwen2.5-Omni-3B last-4 learned mixture + +Recorded on **2026-08-26** from the user-reported remote training and +`infer_turn_head.py` output. This is a single-run diagnostic result, not yet a +multi-seed paper number. + +## Experiment contract + +- Dataset: EasyCom. +- Test split: confirmed as EgoSpeak official Sessions 3, 2, and 1. Training is + recorded as the script's default official Sessions 4--12; verify this against + the remote checkpoint JSON because the supplied excerpt starts at epoch 24. +- Backbone: frozen Qwen2.5-Omni-3B Thinker. +- Streaming mode: open-stream replay with an 8-second causal-memory budget. +- Input sampling: 2 FPS video and raw audio. +- Prediction rate: 5 Hz, with ten future horizons from 0.2 to 2.0 seconds. +- Readout: the query states from the final four Thinker layers, ordered oldest + to newest, followed by a learned global softmax layer mixture and the same + `10 x 3` anticipation MLP. +- Training: head and four scalar layer logits only; Thinker remains frozen. +- Evaluator: EgoSpeak full-timeline per-frame AP/mAP. + +The remote artifact is named +`easycom_frozen_3b_8s_last4_learned_mix_50ep.pth`, but the supplied log and the +user report both show that the actual run ended at **epoch 25/25**. The +`50ep` suffix is therefore a naming error and must not be used as provenance. + +## Summary + +All values below are percentages. The comparison baseline is the previously +recorded frozen 3B, 8-second, open-stream, single-final-layer run evaluated on +the same official test sessions. + +| Readout | 0.2s full mAP | Full mean mAP | Delta at 0.2s | Delta mean | +|---|---:|---:|---:|---:| +| Single final-layer query | 71.99768 | 58.51937 | -- | -- | +| **Last-4 learned mixture** | **73.45085** | **60.53423** | **+1.45317** | **+2.01486** | + +The last-4 readout improves every evaluated horizon. The largest gain is +`+2.41084` at 1.0 seconds; the smallest is still a positive `+1.45317` at +0.2 seconds. + +## Per-horizon full-timeline results + +| Horizon | Last-4 mAP | Single-layer mAP | Delta | Background AP | Wearer AP | Others AP | +|---:|---:|---:|---:|---:|---:|---:| +| 0.2s | 73.45085 | 71.99768 | +1.45317 | 56.69385 | 78.35111 | 85.30760 | +| 0.4s | 68.58645 | 66.74289 | +1.84356 | 49.77891 | 73.48058 | 82.49988 | +| 0.6s | 64.44213 | 62.38972 | +2.05241 | 44.51563 | 68.39317 | 80.41758 | +| 0.8s | 61.34852 | 59.12229 | +2.22623 | 40.98328 | 64.20310 | 78.85917 | +| 1.0s | 59.12464 | 56.71380 | +2.41084 | 38.99833 | 60.75541 | 77.62018 | +| 1.2s | 57.53667 | 55.36300 | +2.17367 | 37.76815 | 58.15110 | 76.69076 | +| 1.4s | 56.43710 | 54.21297 | +2.22413 | 37.22701 | 56.12642 | 75.95786 | +| 1.6s | 55.42035 | 53.43221 | +1.98814 | 36.72606 | 54.14045 | 75.39454 | +| 1.8s | 54.76334 | 52.84703 | +1.91631 | 36.45931 | 52.80226 | 75.02845 | +| 2.0s | 54.23227 | 52.37207 | +1.86020 | 36.45292 | 51.56488 | 74.67902 | +| **Mean** | **60.53423** | **58.51937** | **+2.01486** | **41.56035** | **61.79685** | **78.24550** | + +Averaged over horizons, the AP gains relative to the single-layer baseline are +`+1.92941` for background, `+2.88478` for camera-wearer speaking, and +`+1.23041` for other-people speaking. The wearer class contributes the largest +improvement, which is encouraging for the short-horizon turn-taking objective. + +## Learned layer mixture + +| Layer order | Weight | +|---|---:| +| Oldest of selected four | 0.246384 | +| Second | 0.248327 | +| Third | 0.241774 | +| Final layer | 0.263515 | + +The weights remain close to uniform: normalized entropy is `0.999619`, the L1 +distance from `[0.25, 0.25, 0.25, 0.25]` is only `0.02703`, and the largest +weight exceeds the smallest by only `0.021741`. Consequently, this experiment +supports **multi-layer query readout**, but does not yet show that learning the +mixture is better than a fixed last-4 mean. + +## Interpretation and next controls + +1. This is the strongest recorded frozen-Omni EasyCom result in the project so + far: `73.45085` at 0.2 seconds and `60.53423` mean mAP. +2. The gain occurs at all horizons and all three classes, rather than coming + from a single AP outlier. +3. Because the learned weights are nearly uniform, the first required ablation + is a **fixed uniform last-4 mean** with the same feature files, head recipe, + epoch count, and seed. +4. Then compare each of the four individual query layers. This separates layer + ensembling from access to a specifically better intermediate layer. +5. Repeat the single-layer, fixed-mean, and learned-mixture variants with at + least three identical seeds before using the gain as a paper result. +6. The supplied output contains no strict causal-valid table. Archive the + remote `metrics.json` or rerun the evaluator before the publication table. +7. The single-layer comparison came from an earlier run; its random seed and + head-training duration must be matched in the fixed-mean ablation. + +## Provenance + +- Remote feature root: + `/data1/yfzx/wangx/Datasets/easycom/omni_features/open_stream_2fps_8s_last4_stack` +- Remote checkpoint: + `checkpoints/easycom_frozen_3b_8s_last4_learned_mix_50ep.pth` +- Remote result directory: + `results/easycom_frozen_3b_8s_last4_learned_mix_50ep` +- Final training loss: `0.822849`. +- Raw console excerpt: [`raw_console.txt`](raw_console.txt). +- Machine-readable values: [`summary.json`](summary.json) and + [`metrics.csv`](metrics.csv). +- Raw remote `metrics.json` archived locally: no. +- Exact training-session list verified from the remote checkpoint JSON: no. +- Multi-seed: no. diff --git a/results/easycom_frozen_3b_8s_last4_learned_mix_2026-08-26/metrics.csv b/results/easycom_frozen_3b_8s_last4_learned_mix_2026-08-26/metrics.csv new file mode 100644 index 0000000..0dff93d --- /dev/null +++ b/results/easycom_frozen_3b_8s_last4_learned_mix_2026-08-26/metrics.csv @@ -0,0 +1,12 @@ +horizon_seconds,last4_map,single_final_layer_map,delta,background_ap,camera_wearer_ap,other_people_ap +0.2,73.45085,71.99768,1.45317,56.69385,78.35111,85.30760 +0.4,68.58645,66.74289,1.84356,49.77891,73.48058,82.49988 +0.6,64.44213,62.38972,2.05241,44.51563,68.39317,80.41758 +0.8,61.34852,59.12229,2.22623,40.98328,64.20310,78.85917 +1.0,59.12464,56.71380,2.41084,38.99833,60.75541,77.62018 +1.2,57.53667,55.36300,2.17367,37.76815,58.15110,76.69076 +1.4,56.43710,54.21297,2.22413,37.22701,56.12642,75.95786 +1.6,55.42035,53.43221,1.98814,36.72606,54.14045,75.39454 +1.8,54.76334,52.84703,1.91631,36.45931,52.80226,75.02845 +2.0,54.23227,52.37207,1.86020,36.45292,51.56488,74.67902 +mean,60.53423,58.51937,2.01486,41.56035,61.79685,78.24550 diff --git a/results/easycom_frozen_3b_8s_last4_learned_mix_2026-08-26/raw_console.txt b/results/easycom_frozen_3b_8s_last4_learned_mix_2026-08-26/raw_console.txt new file mode 100644 index 0000000..f836612 --- /dev/null +++ b/results/easycom_frozen_3b_8s_last4_learned_mix_2026-08-26/raw_console.txt @@ -0,0 +1,33 @@ +head epoch 24/25: 100% | 142/142 | 137.53 batch/s +epoch 24: train_loss=0.822925, layer_weights=[0.2464,0.2483,0.2418,0.2635] +head epoch 25/25: 100% | 142/142 | 139.49 batch/s, lr=0 +epoch 25: train_loss=0.822849, layer_weights=[0.2464,0.2483,0.2418,0.2635] +checkpoint: checkpoints/easycom_frozen_3b_8s_last4_learned_mix_50ep.pth +training summary: checkpoints/easycom_frozen_3b_8s_last4_learned_mix_50ep.pth.json +saved layer mixture weights (oldest->newest): 0.246384, 0.248327, 0.241774, 0.263515 + +feature root: /data1/yfzx/wangx/Datasets/easycom/omni_features/open_stream_2fps_8s_last4_stack +test sessions: Session_3, Session_2, Session_1 + +egospeak_full_timeline: +Action anticipation (0.20s) perframe mAP: 73.45085 + AP: background=56.69385, camera_wearer_speaking=78.35111, other_people_speaking=85.30760; n=23048 +Action anticipation (0.40s) perframe mAP: 68.58645 + AP: background=49.77891, camera_wearer_speaking=73.48058, other_people_speaking=82.49988; n=23048 +Action anticipation (0.60s) perframe mAP: 64.44213 + AP: background=44.51563, camera_wearer_speaking=68.39317, other_people_speaking=80.41758; n=23048 +Action anticipation (0.80s) perframe mAP: 61.34852 + AP: background=40.98328, camera_wearer_speaking=64.20310, other_people_speaking=78.85917; n=23048 +Action anticipation (1.00s) perframe mAP: 59.12464 + AP: background=38.99833, camera_wearer_speaking=60.75541, other_people_speaking=77.62018; n=23048 +Action anticipation (1.20s) perframe mAP: 57.53667 + AP: background=37.76815, camera_wearer_speaking=58.15110, other_people_speaking=76.69076; n=23048 +Action anticipation (1.40s) perframe mAP: 56.43710 + AP: background=37.22701, camera_wearer_speaking=56.12642, other_people_speaking=75.95786; n=23048 +Action anticipation (1.60s) perframe mAP: 55.42035 + AP: background=36.72606, camera_wearer_speaking=54.14045, other_people_speaking=75.39454; n=23048 +Action anticipation (1.80s) perframe mAP: 54.76334 + AP: background=36.45931, camera_wearer_speaking=52.80226, other_people_speaking=75.02845; n=23048 +Action anticipation (2.00s) perframe mAP: 54.23227 + AP: background=36.45292, camera_wearer_speaking=51.56488, other_people_speaking=74.67902; n=23048 +Action anticipation (mean) perframe mAP: 60.53423 diff --git a/results/easycom_frozen_3b_8s_last4_learned_mix_2026-08-26/summary.json b/results/easycom_frozen_3b_8s_last4_learned_mix_2026-08-26/summary.json new file mode 100644 index 0000000..ce2357f --- /dev/null +++ b/results/easycom_frozen_3b_8s_last4_learned_mix_2026-08-26/summary.json @@ -0,0 +1,66 @@ +{ + "schema_version": 1, + "recorded_date": "2026-08-26", + "status": "single_run_diagnostic_snapshot", + "dataset": "EasyCom", + "split": { + "train_sessions": [4, 5, 6, 7, 8, 9, 10, 11, 12], + "train_sessions_source": "train_turn_head.py default; not visible in supplied epoch-24 excerpt", + "test_sessions": [3, 2, 1], + "test_sessions_confirmed_in_console": true + }, + "model": { + "backbone": "Qwen2.5-Omni-3B", + "thinker_frozen": true, + "streaming_mode": "open-stream-replay", + "memory_seconds": 8, + "video_fps": 2, + "prediction_hz": 5 + }, + "readout": { + "selected_layers": "last_4_query_states", + "layer_order": "oldest_to_newest", + "fusion": "learned_global_softmax", + "weights": [0.246384, 0.248327, 0.241774, 0.263515], + "normalized_entropy": 0.999619394230947, + "l1_distance_from_uniform": 0.02703, + "max_min_weight_gap": 0.021741 + }, + "training": { + "actual_epochs": 25, + "final_train_loss": 0.822849, + "artifact_name": "easycom_frozen_3b_8s_last4_learned_mix_50ep.pth", + "artifact_epoch_suffix_is_incorrect": true + }, + "metrics_percent": { + "protocol": "egospeak_full_timeline", + "horizon_seconds": [0.2, 0.4, 0.6, 0.8, 1.0, 1.2, 1.4, 1.6, 1.8, 2.0], + "map": [73.45085, 68.58645, 64.44213, 61.34852, 59.12464, 57.53667, 56.43710, 55.42035, 54.76334, 54.23227], + "background_ap": [56.69385, 49.77891, 44.51563, 40.98328, 38.99833, 37.76815, 37.22701, 36.72606, 36.45931, 36.45292], + "camera_wearer_ap": [78.35111, 73.48058, 68.39317, 64.20310, 60.75541, 58.15110, 56.12642, 54.14045, 52.80226, 51.56488], + "other_people_ap": [85.30760, 82.49988, 80.41758, 78.85917, 77.62018, 76.69076, 75.95786, 75.39454, 75.02845, 74.67902], + "mean_map": 60.53423, + "strict_causal_valid_reported": false + }, + "comparison_to_single_final_layer_8s_percent": { + "baseline_mean_map": 58.51937, + "baseline_horizon_0.2_map": 71.99768, + "mean_map_delta": 2.01486, + "horizon_0.2_map_delta": 1.45317, + "horizon_map_delta": [1.45317, 1.84356, 2.05241, 2.22623, 2.41084, 2.17367, 2.22413, 1.98814, 1.91631, 1.86020], + "mean_class_ap_delta": { + "background": 1.92941, + "camera_wearer": 2.88478, + "other_people": 1.23041 + } + }, + "provenance": { + "source": "user-reported remote console output", + "remote_feature_root": "/data1/yfzx/wangx/Datasets/easycom/omni_features/open_stream_2fps_8s_last4_stack", + "remote_checkpoint": "checkpoints/easycom_frozen_3b_8s_last4_learned_mix_50ep.pth", + "remote_result_dir": "results/easycom_frozen_3b_8s_last4_learned_mix_50ep", + "raw_metrics_json_archived": false, + "remote_checkpoint_json_archived": false, + "multi_seed": false + } +} diff --git a/train_omni_turn_lora.py b/train_omni_turn_lora.py index 3232591..724f2ba 100644 --- a/train_omni_turn_lora.py +++ b/train_omni_turn_lora.py @@ -979,6 +979,9 @@ def _save_weight_checkpoint( "num_horizons": NUM_HORIZONS, "num_classes": NUM_CLASSES, "dropout": raw_head.dropout, + "num_input_layers": raw_head.num_input_layers, + "layer_fusion": raw_head.layer_fusion, + "learnable_layer_mixture": raw_head.learnable_layer_mixture, }, "initial_head_checkpoint": ( str(args.head_checkpoint.resolve()) if args.head_checkpoint else None diff --git a/train_turn_head.py b/train_turn_head.py index 6e472d1..cecebaf 100644 --- a/train_turn_head.py +++ b/train_turn_head.py @@ -23,6 +23,7 @@ from tqdm.auto import tqdm from easycom_head import ( CLASS_NAMES, DATASET_NAMES, + LAYER_FUSION_MODES, NUM_CLASSES, NUM_HORIZONS, OFFICIAL_TRAIN_SESSIONS, @@ -95,8 +96,17 @@ def parse_args() -> argparse.Namespace: "--learnable-layer-mixture", action="store_true", help=( - "Learn softmax weights over stacked [T,L,D] query features before " - "the anticipation head." + "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( @@ -135,6 +145,14 @@ def _validate_args(args: argparse.Namespace) -> None: raise ValueError("--num-workers 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: @@ -348,6 +366,7 @@ def _checkpoint( "num_classes": NUM_CLASSES, "dropout": args.dropout, "num_input_layers": num_input_layers, + "layer_fusion": args.layer_fusion, "learnable_layer_mixture": args.learnable_layer_mixture, }, "alignment": { @@ -383,7 +402,8 @@ def _checkpoint( "scheduler": "linear_warmup_times_cosine", "seed": args.seed, "expected_mode": args.expected_mode, - "learnable_layer_mixture": args.learnable_layer_mixture, + "layer_fusion": args.layer_fusion, + "learnable_layer_mixture": args.layer_fusion == "scalar-mixture", "supervision_policy": ( "transition-sparse-manifest" if args.sparse_anchor_cache_dir @@ -401,6 +421,11 @@ def _checkpoint( 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 @@ -473,14 +498,15 @@ def main() -> None: f"training examples: {len(states)}, feature_dim={hidden_size}, " f"input_layers={num_input_layers}" ) - if num_input_layers > 1 and not args.learnable_layer_mixture: + if num_input_layers > 1 and args.layer_fusion == "single": raise ValueError( f"Features contain {num_input_layers} stacked layers; pass " - "--learnable-layer-mixture" + "--layer-fusion fixed-mean, scalar-mixture, or concat-mlp" ) - if num_input_layers == 1 and args.learnable_layer_mixture: + if num_input_layers == 1 and args.layer_fusion != "single": raise ValueError( - "--learnable-layer-mixture requires stacked [T,L,D] features" + f"--layer-fusion {args.layer_fusion} requires stacked [T,L,D] " + "features" ) dev_states: torch.Tensor | None = None dev_targets: torch.Tensor | None = None @@ -529,7 +555,21 @@ def main() -> None: dropout=args.dropout, num_input_layers=num_input_layers, learnable_layer_mixture=args.learnable_layer_mixture, + layer_fusion=args.layer_fusion, ).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"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: @@ -698,7 +738,7 @@ def main() -> None: 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 args.learnable_layer_mixture: + 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(