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

Add transition-aware R1 and R2 objectives

Browse Source
This commit is contained in:
Xin Wang
2026-08-26 17:18:50 +08:00
parent 0590df1e88
commit 5923b89d60
27 changed files with 979 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
train_turn_head.py
+135 -15
View File
@@ -28,6 +28,7 @@ from easycom_head import (
NUM_HORIZONS,
OFFICIAL_TRAIN_SESSIONS,
STEP_SECONDS,
TRANSITION_OBJECTIVE_NAMES,
TurnPredictionHead,
average_precision_metrics,
egospeak_multihorizon_loss,
@@ -41,6 +42,8 @@ from easycom_head import (
normalize_sample_name,
output_classes_for_loss,
output_space_for_loss,
transition_aware_objective,
transition_targets_from_activity,
validate_disjoint_sessions,
)
@@ -125,6 +128,19 @@ def parse_args() -> argparse.Namespace:
"silence/wearer-only/others-only/overlap softmax logits."
),
)
parser.add_argument(
"--transition-objective",
choices=TRANSITION_OBJECTIVE_NAMES,
default="none",
help=(
"none keeps the activity-only baseline; r1-supervised adds a "
"weighted binary transition head; r2-expected-reward additionally "
"maximizes exact expected transition correctness."
),
)
parser.add_argument("--transition-loss-weight", type=float, default=0.2)
parser.add_argument("--transition-pos-weight", type=float, default=3.0)
parser.add_argument("--transition-reward-weight", type=float, default=0.1)
parser.add_argument("--device", default="cuda:0")
parser.add_argument("--num-workers", type=int, default=0)
parser.add_argument("--seed", type=int, default=0)
@@ -154,6 +170,12 @@ def _validate_args(args: argparse.Namespace) -> None:
raise ValueError("--dropout must be in [0, 1)")
if args.num_workers < 0:
raise ValueError("--num-workers cannot be negative")
if args.transition_loss_weight < 0:
raise ValueError("--transition-loss-weight cannot be negative")
if args.transition_pos_weight <= 0:
raise ValueError("--transition-pos-weight must be positive")
if args.transition_reward_weight < 0:
raise ValueError("--transition-reward-weight cannot be negative")
if args.dataset == "ego4d" and args.manifest is None:
raise ValueError("--manifest is required when --dataset ego4d")
if (
@@ -188,9 +210,17 @@ def _load_examples(
expected_mode: str | None,
dataset: str,
selected_indices: dict[str, set[int]] | None = None,
) -> tuple[torch.Tensor, torch.Tensor, int, int, list[dict[str, Any]]]:
) -> tuple[
torch.Tensor,
torch.Tensor,
torch.Tensor,
int,
int,
list[dict[str, Any]],
]:
all_states: list[torch.Tensor] = []
all_targets: list[torch.Tensor] = []
all_transition_targets: list[torch.Tensor] = []
metadata: list[dict[str, Any]] = []
hidden_size: int | None = None
num_input_layers: int | None = None
@@ -203,18 +233,24 @@ def _load_examples(
dataset=dataset,
)
states, targets, future_indices = data.future_examples(NUM_HORIZONS)
current_indices = future_indices[:, 0] - 1
activity_sequence = torch.cat(
(data.targets[current_indices, None, :], targets),
dim=1,
)
transition_targets = transition_targets_from_activity(activity_sequence)
dense_example_count = len(states)
if selected_indices is not None:
if name not in selected_indices:
raise ValueError(f"No sparse anchor manifest was supplied for {name}")
wanted = selected_indices[name]
current_indices = future_indices[:, 0] - 1
keep = torch.tensor(
[int(index) in wanted for index in current_indices.tolist()],
dtype=torch.bool,
)
states = states[keep]
targets = targets[keep]
transition_targets = transition_targets[keep]
if not len(states):
raise ValueError(f"Sparse anchor selection kept no examples for {name}")
if hidden_size is None:
@@ -233,12 +269,18 @@ def _load_examples(
)
all_states.append(states)
all_targets.append(targets)
all_transition_targets.append(transition_targets)
metadata.append(
{
"session": name,
"num_states": len(data.states),
"num_training_examples": len(states),
"num_dense_training_examples": dense_example_count,
"num_positive_transition_targets": int(
transition_targets.sum().item()
),
"transition_target_count": int(transition_targets.numel()),
"transition_positive_rate": float(transition_targets.mean()),
"supervision_policy": (
"transition-sparse-manifest"
if selected_indices is not None
@@ -257,6 +299,7 @@ def _load_examples(
return (
torch.cat(all_states, dim=0),
torch.cat(all_targets, dim=0),
torch.cat(all_transition_targets, dim=0),
hidden_size,
num_input_layers,
metadata,
@@ -372,7 +415,7 @@ def _checkpoint(
history: list[dict[str, Any]],
) -> dict[str, Any]:
return {
"format_version": 2,
"format_version": 3,
"head": head.state_dict(),
"head_config": {
"hidden_size": hidden_size,
@@ -382,6 +425,7 @@ def _checkpoint(
"num_input_layers": num_input_layers,
"layer_fusion": args.layer_fusion,
"learnable_layer_mixture": args.learnable_layer_mixture,
"transition_auxiliary": head.transition_auxiliary,
},
"alignment": {
"state_label_offset": 0,
@@ -414,6 +458,10 @@ def _checkpoint(
"warmup_epochs": args.warmup_epochs,
"warmup_factor": args.warmup_factor,
"loss": args.loss,
"transition_objective": args.transition_objective,
"transition_loss_weight": args.transition_loss_weight,
"transition_pos_weight": args.transition_pos_weight,
"transition_reward_weight": args.transition_reward_weight,
"optimizer": "Adam",
"scheduler": "linear_warmup_times_cosine",
"seed": args.seed,
@@ -502,7 +550,14 @@ 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, num_input_layers, feature_metadata = _load_examples(
(
states,
targets,
transition_targets,
hidden_size,
num_input_layers,
feature_metadata,
) = _load_examples(
args.feature_root,
args.target_root,
train_sessions,
@@ -514,6 +569,11 @@ def main() -> None:
f"training examples: {len(states)}, feature_dim={hidden_size}, "
f"input_layers={num_input_layers}"
)
print(
"transition targets: "
f"{int(transition_targets.sum().item())}/{transition_targets.numel()} "
f"positive ({float(transition_targets.mean()) * 100:.3f}%)"
)
if num_input_layers > 1 and args.layer_fusion == "single":
raise ValueError(
f"Features contain {num_input_layers} stacked layers; pass "
@@ -530,6 +590,7 @@ def main() -> None:
(
dev_states,
dev_targets,
_,
dev_hidden_size,
dev_num_input_layers,
_,
@@ -553,7 +614,7 @@ def main() -> None:
)
print(f"development examples: {len(dev_states)}")
dataset = TensorDataset(states, targets)
dataset = TensorDataset(states, targets, transition_targets)
generator = torch.Generator().manual_seed(args.seed)
loader = DataLoader(
dataset,
@@ -573,6 +634,7 @@ def main() -> None:
num_input_layers=num_input_layers,
learnable_layer_mixture=args.learnable_layer_mixture,
layer_fusion=args.layer_fusion,
transition_auxiliary=args.transition_objective != "none",
).to(device)
trainable_parameters = sum(
parameter.numel() for parameter in head.parameters() if parameter.requires_grad
@@ -584,7 +646,11 @@ def main() -> None:
)
print(
f"loss={args.loss}, output_space={output_space_for_loss(args.loss)}, "
f"head_outputs={num_output_classes}, layer_fusion={args.layer_fusion}, "
f"transition_objective={args.transition_objective}, "
f"activity_outputs={num_output_classes}, "
"transition_outputs="
f"{NUM_HORIZONS if args.transition_objective != 'none' else 0}, "
f"layer_fusion={args.layer_fusion}, "
f"fusion_parameters={fusion_parameters:,}, "
f"head_trainable_parameters={trainable_parameters:,}"
)
@@ -633,6 +699,9 @@ def main() -> None:
for epoch in range(1, args.epochs + 1):
head.train()
loss_sum = 0.0
activity_loss_sum = 0.0
transition_loss_sum = 0.0
transition_reward_sum = 0.0
example_count = 0
progress = tqdm(
loader,
@@ -640,7 +709,7 @@ def main() -> None:
unit="batch",
disable=args.no_progress,
)
for state_batch, target_batch in progress:
for state_batch, target_batch, transition_target_batch in progress:
state_batch = state_batch.to(
device=device,
dtype=torch.float32,
@@ -651,25 +720,64 @@ def main() -> None:
dtype=torch.float32,
non_blocking=True,
)
transition_target_batch = transition_target_batch.to(
device=device,
dtype=torch.float32,
non_blocking=True,
)
optimizer.zero_grad(set_to_none=True)
logits = head(state_batch)
loss = loss_function(logits, target_batch)
if args.transition_objective == "none":
logits = head(state_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)
activity_loss = loss_function(logits, target_batch)
components = transition_aware_objective(
activity_loss,
transition_logits,
transition_target_batch,
objective=args.transition_objective,
transition_weight=args.transition_loss_weight,
pos_weight=args.transition_pos_weight,
reward_weight=args.transition_reward_weight,
)
loss = components["total"]
transition_loss = components["transition"]
expected_reward = components["expected_reward"]
loss.backward()
optimizer.step()
scheduler.step()
batch_count = len(state_batch)
loss_sum += float(loss.detach()) * batch_count
activity_loss_sum += float(activity_loss.detach()) * batch_count
transition_loss_sum += float(transition_loss.detach()) * batch_count
transition_reward_sum += float(expected_reward.detach()) * batch_count
example_count += batch_count
progress.set_postfix(
loss=f"{float(loss.detach()):.5f}",
lr=f"{scheduler.get_last_lr()[0]:.3g}",
)
postfix = {
"loss": f"{float(loss.detach()):.5f}",
"lr": f"{scheduler.get_last_lr()[0]:.3g}",
}
if args.transition_objective != "none":
postfix["transition"] = f"{float(transition_loss.detach()):.5f}"
postfix["reward"] = f"{float(expected_reward.detach()):.5f}"
progress.set_postfix(postfix)
train_loss = loss_sum / example_count
record: dict[str, Any] = {
"epoch": epoch,
"train_loss": train_loss,
"train_activity_loss": activity_loss_sum / example_count,
"learning_rate": scheduler.get_last_lr()[0],
}
if args.transition_objective != "none":
record["train_transition_loss"] = (
transition_loss_sum / example_count
)
record["train_expected_transition_reward"] = (
transition_reward_sum / example_count
)
layer_weights = head.layer_mixture_weights()
if layer_weights is not None:
record["layer_mixture_weights_oldest_to_newest"] = (
@@ -685,6 +793,16 @@ def main() -> None:
)
+ "]"
)
transition_text = (
""
if args.transition_objective == "none"
else (
f", activity_loss={record['train_activity_loss']:.6f}, "
f"transition_loss={record['train_transition_loss']:.6f}, "
"expected_transition_reward="
f"{record['train_expected_transition_reward']:.6f}"
)
)
if dev_sessions:
assert dev_states is not None and dev_targets is not None
@@ -700,7 +818,8 @@ 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}{weight_text}"
f"dev_valid_mean_mAP={dev_map * 100:.3f}"
f"{transition_text}{weight_text}"
)
if dev_map > float(best_dev_map):
best_dev_map = dev_map
@@ -725,7 +844,8 @@ def main() -> None:
print(f"saved new best checkpoint: {args.output}")
else:
print(
f"epoch {epoch:02d}: train_loss={train_loss:.6f}{weight_text}"
f"epoch {epoch:02d}: train_loss={train_loss:.6f}"
f"{transition_text}{weight_text}"
)
history.append(record)