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

Add learnable last-four-layer query fusion

Browse Source
This commit is contained in:
Xin Wang
2026-08-26 10:29:43 +08:00
parent dbef1fc173
commit 43e0e14cf5
7 changed files with 268 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

107
train_turn_head.py
View File

@@ -91,6 +91,14 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--warmup-epochs", type=float, default=10.0)
parser.add_argument("--warmup-factor", type=float, default=0.3)
parser.add_argument("--dropout", type=float, default=0.2)
parser.add_argument(
"--learnable-layer-mixture",
action="store_true",
help=(
"Learn softmax weights over stacked [T,L,D] query features before "
"the anticipation head."
),
)
parser.add_argument(
"--loss",
choices=("lstr-mce", "miniroad-normalized"),
@@ -151,11 +159,12 @@ def _load_examples(
expected_mode: str | None,
dataset: str,
selected_indices: dict[str, set[int]] | None = None,
) -> tuple[torch.Tensor, torch.Tensor, int, list[dict[str, Any]]]:
) -> tuple[torch.Tensor, torch.Tensor, int, int, list[dict[str, Any]]]:
all_states: list[torch.Tensor] = []
all_targets: list[torch.Tensor] = []
metadata: list[dict[str, Any]] = []
hidden_size: int | None = None
num_input_layers: int | None = None
for name in sessions:
data = load_session_data(
feature_root,
@@ -186,6 +195,13 @@ def _load_examples(
f"Feature dimensions differ: {name} has {data.hidden_size}, "
f"expected {hidden_size}"
)
if num_input_layers is None:
num_input_layers = data.num_input_layers
elif data.num_input_layers != num_input_layers:
raise ValueError(
f"Feature layer counts differ: {name} has "
f"{data.num_input_layers}, expected {num_input_layers}"
)
all_states.append(states)
all_targets.append(targets)
metadata.append(
@@ -207,12 +223,13 @@ def _load_examples(
f"valid_10_horizon_examples={dense_example_count}, "
f"selected_examples={len(states)}"
)
if hidden_size is None:
if hidden_size is None or num_input_layers is None:
raise ValueError("No sessions selected")
return (
torch.cat(all_states, dim=0),
torch.cat(all_targets, dim=0),
hidden_size,
num_input_layers,
metadata,
)
@@ -318,6 +335,7 @@ def _checkpoint(
train_sessions: list[str],
dev_sessions: list[str],
hidden_size: int,
num_input_layers: int,
feature_metadata: list[dict[str, Any]],
history: list[dict[str, Any]],
) -> dict[str, Any]:
@@ -329,6 +347,8 @@ def _checkpoint(
"num_horizons": NUM_HORIZONS,
"num_classes": NUM_CLASSES,
"dropout": args.dropout,
"num_input_layers": num_input_layers,
"learnable_layer_mixture": args.learnable_layer_mixture,
},
"alignment": {
"state_label_offset": 0,
@@ -363,6 +383,7 @@ def _checkpoint(
"scheduler": "linear_warmup_times_cosine",
"seed": args.seed,
"expected_mode": args.expected_mode,
"learnable_layer_mixture": args.learnable_layer_mixture,
"supervision_policy": (
"transition-sparse-manifest"
if args.sparse_anchor_cache_dir
@@ -440,7 +461,7 @@ def main() -> None:
print(f"device: {device}")
print(f"train sessions: {', '.join(train_sessions)}")
print(f"dev sessions: {', '.join(dev_sessions) if dev_sessions else 'none'}")
states, targets, hidden_size, feature_metadata = _load_examples(
states, targets, hidden_size, num_input_layers, feature_metadata = _load_examples(
args.feature_root,
args.target_root,
train_sessions,
@@ -448,11 +469,29 @@ def main() -> None:
args.dataset,
sparse_anchor_indices,
)
print(f"training examples: {len(states)}, feature_dim={hidden_size}")
print(
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:
raise ValueError(
f"Features contain {num_input_layers} stacked layers; pass "
"--learnable-layer-mixture"
)
if num_input_layers == 1 and args.learnable_layer_mixture:
raise ValueError(
"--learnable-layer-mixture requires stacked [T,L,D] features"
)
dev_states: torch.Tensor | None = None
dev_targets: torch.Tensor | None = None
if dev_sessions:
dev_states, dev_targets, dev_hidden_size, _ = _load_examples(
(
dev_states,
dev_targets,
dev_hidden_size,
dev_num_input_layers,
_,
) = _load_examples(
args.feature_root,
args.target_root,
dev_sessions,
@@ -465,6 +504,11 @@ def main() -> None:
f"Dev feature dimension {dev_hidden_size} differs from "
f"training dimension {hidden_size}"
)
if dev_num_input_layers != num_input_layers:
raise ValueError(
f"Dev feature layer count {dev_num_input_layers} differs from "
f"training count {num_input_layers}"
)
print(f"development examples: {len(dev_states)}")
dataset = TensorDataset(states, targets)
@@ -483,9 +527,30 @@ def main() -> None:
num_horizons=NUM_HORIZONS,
num_classes=NUM_CLASSES,
dropout=args.dropout,
num_input_layers=num_input_layers,
learnable_layer_mixture=args.learnable_layer_mixture,
).to(device)
if head.layer_logits is None:
optimizer_parameters: Any = head.parameters()
else:
optimizer_parameters = [
{
"params": [
parameter
for name, parameter in head.named_parameters()
if name != "layer_logits"
],
"weight_decay": args.weight_decay,
},
{
"params": [head.layer_logits],
# Weight decay on softmax logits artificially pulls the layer
# distribution back toward uniform.
"weight_decay": 0.0,
},
]
optimizer = torch.optim.Adam(
head.parameters(),
optimizer_parameters,
lr=args.learning_rate,
weight_decay=args.weight_decay,
)
@@ -546,6 +611,21 @@ def main() -> None:
"train_loss": train_loss,
"learning_rate": scheduler.get_last_lr()[0],
}
layer_weights = head.layer_mixture_weights()
if layer_weights is not None:
record["layer_mixture_weights_oldest_to_newest"] = (
layer_weights.detach().float().cpu().tolist()
)
weight_text = (
""
if layer_weights is None
else ", layer_weights=["
+ ",".join(
f"{value:.4f}"
for value in record["layer_mixture_weights_oldest_to_newest"]
)
+ "]"
)
if dev_sessions:
assert dev_states is not None and dev_targets is not None
@@ -560,7 +640,7 @@ def main() -> None:
record["dev_valid_mean_map"] = dev_map
print(
f"epoch {epoch:02d}: train_loss={train_loss:.6f}, "
f"dev_valid_mean_mAP={dev_map * 100:.3f}"
f"dev_valid_mean_mAP={dev_map * 100:.3f}{weight_text}"
)
if dev_map > float(best_dev_map):
best_dev_map = dev_map
@@ -576,6 +656,7 @@ def main() -> None:
train_sessions=train_sessions,
dev_sessions=dev_sessions,
hidden_size=hidden_size,
num_input_layers=num_input_layers,
feature_metadata=feature_metadata,
history=history + [record],
),
@@ -583,7 +664,9 @@ def main() -> None:
)
print(f"saved new best checkpoint: {args.output}")
else:
print(f"epoch {epoch:02d}: train_loss={train_loss:.6f}")
print(
f"epoch {epoch:02d}: train_loss={train_loss:.6f}{weight_text}"
)
history.append(record)
if not dev_sessions:
@@ -598,6 +681,7 @@ def main() -> None:
train_sessions=train_sessions,
dev_sessions=dev_sessions,
hidden_size=hidden_size,
num_input_layers=num_input_layers,
feature_metadata=feature_metadata,
history=history,
),
@@ -614,6 +698,13 @@ 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:
saved_logits = payload["head"]["layer_logits"].float()
saved_weights = saved_logits.softmax(dim=0).tolist()
print(
"saved layer mixture weights (oldest->newest): "
+ ", ".join(f"{value:.6f}" for value in saved_weights)
)
if __name__ == "__main__":