Add causal person fusion and MuVAP stage-3 baselines

This commit is contained in:
Xin Wang
2026-09-10 06:54:22 +08:00
parent 0534c9d3d0
commit d426b4740a
225 changed files with 27716 additions and 119 deletions
+95 -23
View File
@@ -8,6 +8,7 @@ decay, linear warmup for 10 epochs, cosine decay, dropout 0.2, and MCE loss.
from __future__ import annotations
import argparse
import hashlib
import json
import math
import random
@@ -23,6 +24,7 @@ from tqdm.auto import tqdm
from easycom_head import (
CLASS_NAMES,
DATASET_NAMES,
HIGHRATE_PERSON_AV_FUSION_MODES,
LAYER_FUSION_MODES,
LOCAL_FUSION_MODES,
LOCAL_FUSION_SCOPES,
@@ -30,6 +32,7 @@ from easycom_head import (
NUM_HORIZONS,
OFFICIAL_TRAIN_SESSIONS,
STEP_SECONDS,
TRACK_PERSON_FUSION_MODES,
TRANSITION_OBJECTIVE_NAMES,
TurnPredictionHead,
average_precision_metrics,
@@ -53,6 +56,12 @@ from easycom_head import (
from local_turn_features import FramePersonFeatures
INITIALIZED_PERSON_RESIDUAL_MODES = (
*TRACK_PERSON_FUSION_MODES,
*HIGHRATE_PERSON_AV_FUSION_MODES,
)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--feature-root", type=Path, required=True)
@@ -69,7 +78,15 @@ def parse_args() -> argparse.Namespace:
default="pooled",
help=(
"pooled keeps the existing baseline; frame-set-attention is F1; "
"track-gru-residual adds a zero-gated causal track correction; "
"track-gru-residual preserves the legacy zero-gated screen; "
"causal-person-adapter-v2 adds a dynamic-gated, zero-output "
"causal track correction to a matched pooled AV checkpoint; "
"causal-person-adapter-egospeak-gru keeps that adapter but uses "
"EgoSpeak MiniROADA's Linear-LN-ReLU-Dropout-GRU-ReLU temporal "
"stem on each anonymous person track; "
"highrate-person-av-attention lets global Omni plus one shared "
"traditional CPC audio stream query causal 25 Hz person visual/ASD "
"states sampled at the 5 Hz cutoffs; "
"audio-person-frame-set and gated-person-state-adapter are "
"current-sidecar F5/F6 screening variants, not full F5/F6."
),
@@ -96,8 +113,8 @@ def parse_args() -> argparse.Namespace:
"--initialize-pooled-checkpoint",
type=Path,
help=(
"Required by track-gru-residual: copy the exact pooled G+A+V base "
"before training the zero-gated track correction."
"Required by person residual modes: copy an exact pooled G+A(+V) "
"base before training the person correction."
),
)
parser.add_argument(
@@ -223,15 +240,15 @@ def _validate_args(args: argparse.Namespace) -> None:
"Person-set --local-fusion modes require --local-feature-root and "
"--local-modalities av"
)
if args.local_fusion == "track-gru-residual":
if args.local_fusion in INITIALIZED_PERSON_RESIDUAL_MODES:
if args.initialize_pooled_checkpoint is None:
raise ValueError(
"track-gru-residual requires --initialize-pooled-checkpoint"
f"{args.local_fusion} requires --initialize-pooled-checkpoint"
)
elif args.initialize_pooled_checkpoint is not None or args.train_pooled_base:
raise ValueError(
"--initialize-pooled-checkpoint/--train-pooled-base require "
"--local-fusion track-gru-residual"
"a person-residual --local-fusion mode"
)
if args.epochs <= 0:
raise ValueError("--epochs must be positive")
@@ -511,7 +528,7 @@ def _evaluate_valid_horizons(
person_mask = None
if people is not None:
rows = torch.arange(start, min(start + batch_size, len(states)))
if head.local_fusion == "track-gru-residual":
if head.local_fusion in TRACK_PERSON_FUSION_MODES:
person_batch, person_mask = people.padded_history(
rows,
max_observations=head.person_history_frames,
@@ -635,7 +652,21 @@ def _checkpoint(
if args.initialize_pooled_checkpoint is not None
else None
),
"initialize_pooled_checkpoint_sha256": getattr(
args,
"initialize_pooled_checkpoint_sha256",
None,
),
"initialize_pooled_checkpoint_epoch": getattr(
args,
"initialize_pooled_checkpoint_epoch",
None,
),
"train_pooled_base": args.train_pooled_base,
"pooled_base_frozen": bool(
args.local_fusion in INITIALIZED_PERSON_RESIDUAL_MODES
and not args.train_pooled_base
),
"learnable_layer_mixture": args.layer_fusion == "scalar-mixture",
"supervision_policy": (
"transition-sparse-manifest"
@@ -651,7 +682,7 @@ def _checkpoint(
}
def _initialize_track_residual_from_pooled(
def _initialize_person_residual_from_pooled(
head: TurnPredictionHead,
checkpoint: Path,
*,
@@ -660,7 +691,7 @@ def _initialize_track_residual_from_pooled(
dev_sessions: list[str],
local_contract: dict[str, Any] | None,
) -> None:
"""Load and optionally freeze the exact pooled base inside a track model."""
"""Load and optionally freeze the exact pooled base in a person model."""
path = checkpoint.expanduser().resolve()
if not path.is_file():
@@ -671,18 +702,28 @@ def _initialize_track_residual_from_pooled(
config = checkpoint_head_config(payload)
if config["local_fusion"] != "pooled":
raise ValueError(
"Track residual initialization requires a pooled checkpoint, got "
"Person residual initialization requires a pooled checkpoint, got "
f"{config['local_fusion']!r}: {path}"
)
if checkpoint_loss_name(payload) != args.loss:
raise ValueError("Pooled checkpoint loss differs from the track experiment")
raise ValueError("Pooled checkpoint loss differs from the person experiment")
if payload.get("dataset", args.dataset) != args.dataset:
raise ValueError("Pooled checkpoint dataset differs from the track experiment")
raise ValueError("Pooled checkpoint dataset differs from the person experiment")
if payload.get("train_sessions") != train_sessions:
raise ValueError("Pooled checkpoint train sessions differ from this run")
if payload.get("dev_sessions", []) != dev_sessions:
raise ValueError("Pooled checkpoint dev sessions differ from this run")
if payload.get("local_contract") != local_contract:
if args.local_fusion in HIGHRATE_PERSON_AV_FUSION_MODES:
expected_base_contract = (
{**local_contract, "modalities": "audio"}
if local_contract is not None
else None
)
expected_local_visual_dim = 0
else:
expected_base_contract = local_contract
expected_local_visual_dim = head.local_visual_dim
if payload.get("local_contract") != expected_base_contract:
raise ValueError("Pooled checkpoint local feature contract differs from this run")
expected = {
@@ -691,9 +732,11 @@ def _initialize_track_residual_from_pooled(
"num_classes": head.num_classes,
"num_input_layers": head.num_input_layers,
"layer_fusion": head.layer_fusion,
"local_visual_dim": head.local_visual_dim,
"local_visual_dim": expected_local_visual_dim,
"local_audio_dim": head.local_audio_dim,
"local_projection_dim": head.local_projection_dim,
"dropout": head.dropout,
"transition_auxiliary": head.transition_auxiliary,
}
mismatches = {
key: (config.get(key), value)
@@ -702,12 +745,17 @@ def _initialize_track_residual_from_pooled(
}
if mismatches:
raise ValueError(
f"Pooled checkpoint architecture differs from track base: {mismatches}"
f"Pooled checkpoint architecture differs from person base: {mismatches}"
)
def is_base_parameter(name: str) -> bool:
return name == "layer_logits" or name.startswith(
("layer_fusion_net.", "local_projections.", "net.")
(
"layer_fusion_net.",
"local_projections.",
"net.",
"transition_head.",
)
)
source = payload["head"]
@@ -728,6 +776,13 @@ def _initialize_track_residual_from_pooled(
current[name] = source[name]
head.load_state_dict(current)
digest = hashlib.sha256()
with path.open("rb") as stream:
for chunk in iter(lambda: stream.read(1024 * 1024), b""):
digest.update(chunk)
args.initialize_pooled_checkpoint_sha256 = digest.hexdigest()
args.initialize_pooled_checkpoint_epoch = payload.get("epoch")
if not args.train_pooled_base:
for name, parameter in head.named_parameters():
if is_base_parameter(name):
@@ -932,9 +987,9 @@ def main() -> None:
person_track_max_gap_seconds=args.person_track_max_gap_seconds,
**local_head_dimensions(local_contract),
).to(device)
if args.local_fusion == "track-gru-residual":
if args.local_fusion in INITIALIZED_PERSON_RESIDUAL_MODES:
assert args.initialize_pooled_checkpoint is not None
_initialize_track_residual_from_pooled(
_initialize_person_residual_from_pooled(
head,
args.initialize_pooled_checkpoint,
args=args,
@@ -1023,13 +1078,21 @@ def main() -> None:
best_epoch = args.epochs if not dev_sessions else 0
for epoch in range(1, args.epochs + 1):
head.train()
if args.local_fusion == "track-gru-residual" and not args.train_pooled_base:
if (
args.local_fusion in INITIALIZED_PERSON_RESIDUAL_MODES
and not args.train_pooled_base
):
# Keep dropout/norm behavior in the frozen pooled checkpoint fixed
# while the new track branch learns a residual correction.
# while the new person branch learns a residual correction. This
# applies to both track-history adapters and the high-rate person
# set; otherwise head.train() would silently turn dropout back on
# inside the supposedly frozen matched base.
if head.layer_fusion_net is not None:
head.layer_fusion_net.eval()
head.local_projections.eval()
head.net.eval()
if head.transition_head is not None:
head.transition_head.eval()
loss_sum = 0.0
activity_loss_sum = 0.0
transition_loss_sum = 0.0
@@ -1069,7 +1132,7 @@ def main() -> None:
person_batch = None
person_mask = None
if person_features is not None:
if head.local_fusion == "track-gru-residual":
if head.local_fusion in TRACK_PERSON_FUSION_MODES:
person_batch, person_mask = person_features.padded_history(
example_index_batch,
max_observations=head.person_history_frames,
@@ -1155,6 +1218,9 @@ def main() -> None:
interaction_gates = head.person_interaction_gates()
if interaction_gates:
record["person_interaction_gates"] = interaction_gates
adapter_diagnostics = head.person_adapter_diagnostics()
if adapter_diagnostics:
record["person_adapter_diagnostics"] = adapter_diagnostics
weight_text = (
""
if layer_weights is None
@@ -1181,6 +1247,12 @@ def main() -> None:
else ", person_gates="
+ json.dumps(interaction_gates, sort_keys=True)
)
adapter_text = (
""
if not adapter_diagnostics
else ", person_adapter="
+ json.dumps(adapter_diagnostics, sort_keys=True)
)
if dev_sessions:
assert dev_states is not None and dev_targets is not None
@@ -1199,7 +1271,7 @@ def main() -> None:
print(
f"epoch {epoch:02d}: train_loss={train_loss:.6f}, "
f"dev_valid_mean_mAP={dev_map * 100:.3f}"
f"{transition_text}{weight_text}{gate_text}"
f"{transition_text}{weight_text}{gate_text}{adapter_text}"
)
if dev_map > float(best_dev_map):
best_dev_map = dev_map
@@ -1225,7 +1297,7 @@ def main() -> None:
else:
print(
f"epoch {epoch:02d}: train_loss={train_loss:.6f}"
f"{transition_text}{weight_text}{gate_text}"
f"{transition_text}{weight_text}{gate_text}{adapter_text}"
)
history.append(record)