Add overlap-aware turn prediction objectives

train_turn_head.py

@@ -24,17 +24,23 @@ from easycom_head import (
     CLASS_NAMES,
     DATASET_NAMES,
     LAYER_FUSION_MODES,
-    NUM_CLASSES,
+    LOSS_NAMES,
     NUM_HORIZONS,
     OFFICIAL_TRAIN_SESSIONS,
     STEP_SECONDS,
     TurnPredictionHead,
     average_precision_metrics,
     egospeak_multihorizon_loss,
+    four_state_multihorizon_loss,
+    head_output_names_for_loss,
+    independent_bce_multihorizon_loss,
     load_ego4d_manifest_split,
     load_session_data,
+    marginal_probabilities_from_logits,
     miniroad_multihorizon_loss,
     normalize_sample_name,
+    output_classes_for_loss,
+    output_space_for_loss,
     validate_disjoint_sessions,
 )
 
@@ -111,8 +117,13 @@ def parse_args() -> argparse.Namespace:
     )
     parser.add_argument(
         "--loss",
-        choices=("lstr-mce", "miniroad-normalized"),
+        choices=LOSS_NAMES,
         default="lstr-mce",
+        help=(
+            "lstr-mce/miniroad-normalized use three competing softmax logits; "
+            "independent-bce uses three sigmoid marginals; four-state-ce uses "
+            "silence/wearer-only/others-only/overlap softmax logits."
+        ),
     )
     parser.add_argument("--device", default="cuda:0")
     parser.add_argument("--num-workers", type=int, default=0)
@@ -314,6 +325,7 @@ def _evaluate_valid_horizons(
     targets: torch.Tensor,
     device: torch.device,
     batch_size: int,
+    loss_name: str,
 ) -> dict[str, Any]:
     head.eval()
     chunks: list[torch.Tensor] = []
@@ -323,7 +335,9 @@ def _evaluate_valid_horizons(
             dtype=torch.float32,
             non_blocking=True,
         )
-        chunks.append(head(x).softmax(dim=-1).cpu())
+        chunks.append(
+            marginal_probabilities_from_logits(head(x), loss_name).cpu()
+        )
     probabilities = torch.cat(chunks, dim=0).numpy()
     target_array = targets.numpy()
 
@@ -358,12 +372,12 @@ def _checkpoint(
     history: list[dict[str, Any]],
 ) -> dict[str, Any]:
     return {
-        "format_version": 1,
+        "format_version": 2,
         "head": head.state_dict(),
         "head_config": {
             "hidden_size": hidden_size,
             "num_horizons": NUM_HORIZONS,
-            "num_classes": NUM_CLASSES,
+            "num_classes": head.num_classes,
             "dropout": args.dropout,
             "num_input_layers": num_input_layers,
             "layer_fusion": args.layer_fusion,
@@ -377,6 +391,8 @@ def _checkpoint(
             ],
         },
         "classes": list(CLASS_NAMES),
+        "head_output_names": list(head_output_names_for_loss(args.loss)),
+        "output_space": output_space_for_loss(args.loss),
         "dataset": args.dataset,
         "manifest": (
             str(args.manifest.expanduser().resolve())
@@ -548,10 +564,11 @@ def main() -> None:
         persistent_workers=args.num_workers > 0,
         generator=generator,
     )
+    num_output_classes = output_classes_for_loss(args.loss)
     head = TurnPredictionHead(
         hidden_size=hidden_size,
         num_horizons=NUM_HORIZONS,
-        num_classes=NUM_CLASSES,
+        num_classes=num_output_classes,
         dropout=args.dropout,
         num_input_layers=num_input_layers,
         learnable_layer_mixture=args.learnable_layer_mixture,
@@ -566,7 +583,8 @@ def main() -> None:
         else (head.layer_logits.numel() if head.layer_logits is not None else 0)
     )
     print(
-        f"layer_fusion={args.layer_fusion}, "
+        f"loss={args.loss}, output_space={output_space_for_loss(args.loss)}, "
+        f"head_outputs={num_output_classes}, layer_fusion={args.layer_fusion}, "
         f"fusion_parameters={fusion_parameters:,}, "
         f"head_trainable_parameters={trainable_parameters:,}"
     )
@@ -601,11 +619,12 @@ def main() -> None:
         warmup_epochs=args.warmup_epochs,
         warmup_factor=args.warmup_factor,
     )
-    loss_function = (
-        egospeak_multihorizon_loss
-        if args.loss == "lstr-mce"
-        else miniroad_multihorizon_loss
-    )
+    loss_function = {
+        "lstr-mce": egospeak_multihorizon_loss,
+        "miniroad-normalized": miniroad_multihorizon_loss,
+        "independent-bce": independent_bce_multihorizon_loss,
+        "four-state-ce": four_state_multihorizon_loss,
+    }[args.loss]
 
     args.output.parent.mkdir(parents=True, exist_ok=True)
     history: list[dict[str, Any]] = []
@@ -675,6 +694,7 @@ def main() -> None:
                 dev_targets,
                 device,
                 args.batch_size,
+                args.loss,
             )
             dev_map = float(dev_metrics["mean_map"])
             record["dev_valid_mean_map"] = dev_map