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

Add pooled local AV baselines and research notes

Browse Source
This commit is contained in:
Xin Wang
2026-08-27 14:19:33 +08:00
parent d34725778f
commit 11f197724f
23 changed files with 1541 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
train_turn_head.py
+69 -7
View File
@@ -53,6 +53,15 @@ def parse_args() -> argparse.Namespace:
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(
"--local-feature-root", type=Path,
help="Optional frozen local-feature sidecars; omit for the original baseline.",
)
parser.add_argument("--local-modalities", choices=("visual", "audio", "av"), default="av")
parser.add_argument(
"--local-projection-dim", type=int, default=256,
help="Trainable projection per local modality; 0 retains encoder dimensions.",
)
parser.add_argument(
"--dataset",
choices=DATASET_NAMES,
@@ -154,6 +163,8 @@ def parse_args() -> argparse.Namespace:
def _validate_args(args: argparse.Namespace) -> None:
if args.local_projection_dim < 0:
raise ValueError("--local-projection-dim must be non-negative")
if args.epochs <= 0:
raise ValueError("--epochs must be positive")
if args.batch_size <= 0:
@@ -210,6 +221,9 @@ def _load_examples(
expected_mode: str | None,
dataset: str,
selected_indices: dict[str, set[int]] | None = None,
local_feature_root: Path | None = None,
local_modalities: str = "av",
local_contract: dict[str, Any] | None = None,
) -> tuple[
torch.Tensor,
torch.Tensor,
@@ -217,10 +231,12 @@ def _load_examples(
int,
int,
list[dict[str, Any]],
torch.Tensor,
]:
all_states: list[torch.Tensor] = []
all_targets: list[torch.Tensor] = []
all_transition_targets: list[torch.Tensor] = []
all_local: list[torch.Tensor] = []
metadata: list[dict[str, Any]] = []
hidden_size: int | None = None
num_input_layers: int | None = None
@@ -231,9 +247,18 @@ def _load_examples(
name,
expected_mode=expected_mode,
dataset=dataset,
local_feature_root=local_feature_root,
local_modalities=local_modalities,
local_contract=local_contract,
)
local_contract = data.feature_metadata.get("local_contract")
states, targets, future_indices = data.future_examples(NUM_HORIZONS)
current_indices = future_indices[:, 0] - 1
local = (
data.local_features[current_indices]
if data.local_features is not None
else torch.empty(len(states), 0)
)
activity_sequence = torch.cat(
(data.targets[current_indices, None, :], targets),
dim=1,
@@ -251,6 +276,7 @@ def _load_examples(
states = states[keep]
targets = targets[keep]
transition_targets = transition_targets[keep]
local = local[keep]
if not len(states):
raise ValueError(f"Sparse anchor selection kept no examples for {name}")
if hidden_size is None:
@@ -270,6 +296,7 @@ def _load_examples(
all_states.append(states)
all_targets.append(targets)
all_transition_targets.append(transition_targets)
all_local.append(local)
metadata.append(
{
"session": name,
@@ -303,6 +330,7 @@ def _load_examples(
hidden_size,
num_input_layers,
metadata,
torch.cat(all_local, dim=0),
)
@@ -369,6 +397,7 @@ def _evaluate_valid_horizons(
device: torch.device,
batch_size: int,
loss_name: str,
local: torch.Tensor | None = None,
) -> dict[str, Any]:
head.eval()
chunks: list[torch.Tensor] = []
@@ -378,9 +407,11 @@ def _evaluate_valid_horizons(
dtype=torch.float32,
non_blocking=True,
)
chunks.append(
marginal_probabilities_from_logits(head(x), loss_name).cpu()
side = (
local[start : start + batch_size].to(device=device, dtype=torch.float32)
if local is not None else None
)
chunks.append(marginal_probabilities_from_logits(head(x, side), loss_name).cpu())
probabilities = torch.cat(chunks, dim=0).numpy()
target_array = targets.numpy()
@@ -415,7 +446,7 @@ def _checkpoint(
history: list[dict[str, Any]],
) -> dict[str, Any]:
return {
"format_version": 3,
"format_version": 4,
"head": head.state_dict(),
"head_config": {
"hidden_size": hidden_size,
@@ -426,7 +457,11 @@ def _checkpoint(
"layer_fusion": args.layer_fusion,
"learnable_layer_mixture": args.learnable_layer_mixture,
"transition_auxiliary": head.transition_auxiliary,
"local_visual_dim": head.local_visual_dim,
"local_audio_dim": head.local_audio_dim,
"local_projection_dim": head.local_projection_dim,
},
"local_contract": feature_metadata[0].get("local_contract"),
"alignment": {
"state_label_offset": 0,
"future_offsets": list(range(1, NUM_HORIZONS + 1)),
@@ -467,6 +502,9 @@ def _checkpoint(
"seed": args.seed,
"expected_mode": args.expected_mode,
"layer_fusion": args.layer_fusion,
"local_feature_root": str(args.local_feature_root) if args.local_feature_root else None,
"local_modalities": args.local_modalities if args.local_feature_root else None,
"local_projection_dim": args.local_projection_dim,
"learnable_layer_mixture": args.layer_fusion == "scalar-mixture",
"supervision_policy": (
"transition-sparse-manifest"
@@ -557,6 +595,7 @@ def main() -> None:
hidden_size,
num_input_layers,
feature_metadata,
local_features,
) = _load_examples(
args.feature_root,
args.target_root,
@@ -564,7 +603,12 @@ def main() -> None:
expected_mode,
args.dataset,
sparse_anchor_indices,
args.local_feature_root,
args.local_modalities,
)
local_contract = feature_metadata[0].get("local_contract")
if not len(states):
raise ValueError("No valid training examples")
print(
f"training examples: {len(states)}, feature_dim={hidden_size}, "
f"input_layers={num_input_layers}"
@@ -586,6 +630,7 @@ def main() -> None:
)
dev_states: torch.Tensor | None = None
dev_targets: torch.Tensor | None = None
dev_local: torch.Tensor | None = None
if dev_sessions:
(
dev_states,
@@ -594,6 +639,7 @@ def main() -> None:
dev_hidden_size,
dev_num_input_layers,
_,
dev_local,
) = _load_examples(
args.feature_root,
args.target_root,
@@ -601,6 +647,9 @@ def main() -> None:
expected_mode,
args.dataset,
None,
args.local_feature_root,
args.local_modalities,
local_contract,
)
if dev_hidden_size != hidden_size:
raise ValueError(
@@ -614,7 +663,7 @@ def main() -> None:
)
print(f"development examples: {len(dev_states)}")
dataset = TensorDataset(states, targets, transition_targets)
dataset = TensorDataset(states, targets, transition_targets, local_features)
generator = torch.Generator().manual_seed(args.seed)
loader = DataLoader(
dataset,
@@ -626,6 +675,8 @@ def main() -> None:
generator=generator,
)
num_output_classes = output_classes_for_loss(args.loss)
from local_turn_features import local_head_dimensions
head = TurnPredictionHead(
hidden_size=hidden_size,
num_horizons=NUM_HORIZONS,
@@ -635,7 +686,14 @@ def main() -> None:
learnable_layer_mixture=args.learnable_layer_mixture,
layer_fusion=args.layer_fusion,
transition_auxiliary=args.transition_objective != "none",
local_projection_dim=args.local_projection_dim,
**local_head_dimensions(local_contract),
).to(device)
if local_contract is not None:
print(
f"local fusion: {args.local_modalities}; raw V={head.local_visual_dim}, "
f"A={head.local_audio_dim}; projection={args.local_projection_dim or 'identity'}"
)
trainable_parameters = sum(
parameter.numel() for parameter in head.parameters() if parameter.requires_grad
)
@@ -709,7 +767,10 @@ def main() -> None:
unit="batch",
disable=args.no_progress,
)
for state_batch, target_batch, transition_target_batch in progress:
for state_batch, target_batch, transition_target_batch, local_batch in progress:
local_batch = local_batch.to(
device=device, dtype=torch.float32, non_blocking=True,
)
state_batch = state_batch.to(
device=device,
dtype=torch.float32,
@@ -727,13 +788,13 @@ def main() -> None:
)
optimizer.zero_grad(set_to_none=True)
if args.transition_objective == "none":
logits = head(state_batch)
logits = head(state_batch, local_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)
logits, transition_logits = head.forward_with_transition(state_batch, local_batch)
activity_loss = loss_function(logits, target_batch)
components = transition_aware_objective(
activity_loss,
@@ -813,6 +874,7 @@ def main() -> None:
device,
args.batch_size,
args.loss,
dev_local,
)
dev_map = float(dev_metrics["mean_map"])
record["dev_valid_mean_map"] = dev_map