wx44wx/icassp2027-experiment
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add causal person track GRU residual experiment

Browse Source
This commit is contained in:
Xin Wang
2026-09-07 00:30:23 +08:00
parent 9fc4577aea
commit 4aed4697a5
15 changed files with 1185 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files

185
train_turn_head.py
View File

@@ -33,6 +33,8 @@ from easycom_head import (
TRANSITION_OBJECTIVE_NAMES,
TurnPredictionHead,
average_precision_metrics,
checkpoint_head_config,
checkpoint_loss_name,
egospeak_multihorizon_loss,
four_state_multihorizon_loss,
head_output_names_for_loss,
@@ -67,6 +69,7 @@ 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; "
"audio-person-frame-set and gated-person-state-adapter are "
"current-sidecar F5/F6 screening variants, not full F5/F6."
),
@@ -77,6 +80,31 @@ def parse_args() -> argparse.Namespace:
)
parser.add_argument("--person-attention-dim", type=int, default=256)
parser.add_argument("--person-attention-heads", type=int, default=4)
parser.add_argument(
"--person-history-frames",
type=int,
default=16,
help="Maximum causal observations per current face track (16 at 2 FPS = 8s).",
)
parser.add_argument(
"--person-track-max-gap-seconds",
type=float,
default=1.0,
help="Reset a provided track when consecutive observations are farther apart.",
)
parser.add_argument(
"--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."
),
)
parser.add_argument(
"--train-pooled-base",
action="store_true",
help="Also update the initialized pooled base; default keeps it frozen.",
)
parser.add_argument(
"--dataset",
choices=DATASET_NAMES,
@@ -184,6 +212,10 @@ def _validate_args(args: argparse.Namespace) -> None:
raise ValueError("Person attention dimension/heads must be positive")
if args.person_attention_dim % args.person_attention_heads:
raise ValueError("--person-attention-dim must be divisible by --person-attention-heads")
if args.person_history_frames <= 0:
raise ValueError("--person-history-frames must be positive")
if args.person_track_max_gap_seconds <= 0:
raise ValueError("--person-track-max-gap-seconds must be positive")
if args.local_fusion != "pooled" and (
args.local_feature_root is None or args.local_modalities != "av"
):
@@ -191,6 +223,16 @@ 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.initialize_pooled_checkpoint is None:
raise ValueError(
"track-gru-residual 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"
)
if args.epochs <= 0:
raise ValueError("--epochs must be positive")
if args.batch_size <= 0:
@@ -252,6 +294,7 @@ def _load_examples(
local_contract: dict[str, Any] | None = None,
compute_transition_targets: bool = True,
local_fusion: str = "pooled",
person_track_max_gap_seconds: float = 1.0,
) -> tuple[
torch.Tensor,
torch.Tensor,
@@ -281,6 +324,7 @@ def _load_examples(
local_modalities=local_modalities,
local_contract=local_contract,
local_fusion=local_fusion,
person_track_max_gap_seconds=person_track_max_gap_seconds,
)
local_contract = data.feature_metadata.get("local_contract")
states, targets, future_indices = data.future_examples(NUM_HORIZONS)
@@ -467,7 +511,13 @@ def _evaluate_valid_horizons(
person_mask = None
if people is not None:
rows = torch.arange(start, min(start + batch_size, len(states)))
person_batch, person_mask = people.padded(rows)
if head.local_fusion == "track-gru-residual":
person_batch, person_mask = people.padded_history(
rows,
max_observations=head.person_history_frames,
)
else:
person_batch, person_mask = people.padded(rows)
person_batch = person_batch.to(device=device, dtype=torch.float32)
person_mask = person_mask.to(device=device)
chunks.append(
@@ -510,7 +560,7 @@ def _checkpoint(
history: list[dict[str, Any]],
) -> dict[str, Any]:
return {
"format_version": 6,
"format_version": 7,
"head": head.state_dict(),
"head_config": {
"hidden_size": hidden_size,
@@ -527,6 +577,8 @@ def _checkpoint(
"local_fusion": head.local_fusion,
"person_attention_dim": head.person_attention_dim,
"person_attention_heads": head.person_attention_heads,
"person_history_frames": head.person_history_frames,
"person_track_max_gap_seconds": head.person_track_max_gap_seconds,
},
"local_contract": feature_metadata[0].get("local_contract"),
"alignment": {
@@ -576,6 +628,14 @@ def _checkpoint(
"local_projection_dim": args.local_projection_dim,
"person_attention_dim": args.person_attention_dim,
"person_attention_heads": args.person_attention_heads,
"person_history_frames": args.person_history_frames,
"person_track_max_gap_seconds": args.person_track_max_gap_seconds,
"initialize_pooled_checkpoint": (
str(args.initialize_pooled_checkpoint.expanduser().resolve())
if args.initialize_pooled_checkpoint is not None
else None
),
"train_pooled_base": args.train_pooled_base,
"learnable_layer_mixture": args.layer_fusion == "scalar-mixture",
"supervision_policy": (
"transition-sparse-manifest"
@@ -591,6 +651,94 @@ def _checkpoint(
}
def _initialize_track_residual_from_pooled(
head: TurnPredictionHead,
checkpoint: Path,
*,
args: argparse.Namespace,
train_sessions: list[str],
dev_sessions: list[str],
local_contract: dict[str, Any] | None,
) -> None:
"""Load and optionally freeze the exact pooled base inside a track model."""
path = checkpoint.expanduser().resolve()
if not path.is_file():
raise FileNotFoundError(path)
payload = torch.load(path, map_location="cpu", weights_only=False)
if not isinstance(payload, dict) or not isinstance(payload.get("head"), dict):
raise ValueError(f"Invalid pooled checkpoint: {path}")
config = checkpoint_head_config(payload)
if config["local_fusion"] != "pooled":
raise ValueError(
"Track 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")
if payload.get("dataset", args.dataset) != args.dataset:
raise ValueError("Pooled checkpoint dataset differs from the track 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:
raise ValueError("Pooled checkpoint local feature contract differs from this run")
expected = {
"hidden_size": head.hidden_size,
"num_horizons": head.num_horizons,
"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_audio_dim": head.local_audio_dim,
"local_projection_dim": head.local_projection_dim,
}
mismatches = {
key: (config.get(key), value)
for key, value in expected.items()
if config.get(key) != value
}
if mismatches:
raise ValueError(
f"Pooled checkpoint architecture differs from track base: {mismatches}"
)
def is_base_parameter(name: str) -> bool:
return name == "layer_logits" or name.startswith(
("layer_fusion_net.", "local_projections.", "net.")
)
source = payload["head"]
current = head.state_dict()
required = [name for name in current if is_base_parameter(name)]
missing = [name for name in required if name not in source]
bad_shapes = [
name
for name in required
if name in source and source[name].shape != current[name].shape
]
if missing or bad_shapes:
raise ValueError(
"Pooled checkpoint cannot initialize the exact base; "
f"missing={missing[:8]}, shape_mismatch={bad_shapes[:8]}"
)
for name in required:
current[name] = source[name]
head.load_state_dict(current)
if not args.train_pooled_base:
for name, parameter in head.named_parameters():
if is_base_parameter(name):
parameter.requires_grad_(False)
print(
f"initialized pooled base: {path}; "
f"frozen={not args.train_pooled_base}",
flush=True,
)
def main() -> None:
args = parse_args()
_validate_args(args)
@@ -679,6 +827,7 @@ def main() -> None:
args.local_modalities,
compute_transition_targets=args.transition_objective != "none",
local_fusion=args.local_fusion,
person_track_max_gap_seconds=args.person_track_max_gap_seconds,
)
local_contract = feature_metadata[0].get("local_contract")
if not len(states):
@@ -732,6 +881,7 @@ def main() -> None:
local_contract,
compute_transition_targets=args.transition_objective != "none",
local_fusion=args.local_fusion,
person_track_max_gap_seconds=args.person_track_max_gap_seconds,
)
if dev_hidden_size != hidden_size:
raise ValueError(
@@ -778,8 +928,20 @@ def main() -> None:
local_fusion=args.local_fusion,
person_attention_dim=args.person_attention_dim,
person_attention_heads=args.person_attention_heads,
person_history_frames=args.person_history_frames,
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":
assert args.initialize_pooled_checkpoint is not None
_initialize_track_residual_from_pooled(
head,
args.initialize_pooled_checkpoint,
args=args,
train_sessions=train_sessions,
dev_sessions=dev_sessions,
local_contract=local_contract,
)
if local_contract is not None:
print(
f"local fusion: {args.local_fusion}/{args.local_modalities}; "
@@ -861,6 +1023,13 @@ 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:
# Keep dropout/norm behavior in the frozen pooled checkpoint fixed
# while the new track branch learns a residual correction.
if head.layer_fusion_net is not None:
head.layer_fusion_net.eval()
head.local_projections.eval()
head.net.eval()
loss_sum = 0.0
activity_loss_sum = 0.0
transition_loss_sum = 0.0
@@ -900,9 +1069,15 @@ def main() -> None:
person_batch = None
person_mask = None
if person_features is not None:
person_batch, person_mask = person_features.padded(
example_index_batch
)
if head.local_fusion == "track-gru-residual":
person_batch, person_mask = person_features.padded_history(
example_index_batch,
max_observations=head.person_history_frames,
)
else:
person_batch, person_mask = person_features.padded(
example_index_batch
)
person_batch = person_batch.to(
device=device,
dtype=torch.float32,