Add causal person fusion and MuVAP stage-3 baselines
This commit is contained in:
+95
-23
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user