Add overlap-aware turn prediction objectives

infer_turn_head.py

@@ -21,9 +21,13 @@ from easycom_head import (
     TurnPredictionHead,
     average_precision_metrics,
     checkpoint_head_config,
+    checkpoint_loss_name,
     load_ego4d_manifest_split,
     load_session_data,
+    marginal_probabilities_from_logits,
     normalize_sample_name,
+    output_classes_for_loss,
+    output_space_for_loss,
 )
 
 
@@ -86,9 +90,12 @@ def _load_head(
         raise ValueError(
             f"Checkpoint has {config['num_horizons']} horizons, expected {NUM_HORIZONS}"
         )
-    if int(config["num_classes"]) != NUM_CLASSES:
+    loss_name = checkpoint_loss_name(payload)
+    expected_classes = output_classes_for_loss(loss_name)
+    if int(config["num_classes"]) != expected_classes:
         raise ValueError(
-            f"Checkpoint has {config['num_classes']} classes, expected {NUM_CLASSES}"
+            f"Checkpoint has {config['num_classes']} outputs, expected "
+            f"{expected_classes} for loss {loss_name!r}"
         )
     head = TurnPredictionHead(**config)
     head.load_state_dict(state_dict)
@@ -104,6 +111,7 @@ def _predict_states(
     *,
     description: str,
     show_progress: bool,
+    loss_name: str,
 ) -> np.ndarray:
     chunks: list[torch.Tensor] = []
     starts = range(0, len(states), batch_size)
@@ -120,7 +128,9 @@ def _predict_states(
             dtype=torch.float32,
             non_blocking=True,
         )
-        chunks.append(head(batch).softmax(dim=-1).cpu())
+        chunks.append(
+            marginal_probabilities_from_logits(head(batch), loss_name).cpu()
+        )
     return torch.cat(chunks, dim=0).numpy()
 
 
@@ -258,6 +268,7 @@ def main() -> None:
             f"{args.dataset!r}"
         )
     head_config = checkpoint_head_config(checkpoint_payload)
+    loss_name = checkpoint_loss_name(checkpoint_payload)
 
     args.output_dir.mkdir(parents=True, exist_ok=True)
     prediction_dir = args.output_dir / "pred"
@@ -269,6 +280,9 @@ def main() -> None:
     session_targets: dict[str, np.ndarray] = {}
 
     print(f"device: {device}")
+    print(
+        f"loss: {loss_name}; output_space: {output_space_for_loss(loss_name)}"
+    )
     print(f"sessions: {', '.join(sessions)}")
     for name in sessions:
         data = load_session_data(
@@ -295,6 +309,7 @@ def main() -> None:
             args.batch_size,
             description=f"infer {name}",
             show_progress=not args.no_progress,
+            loss_name=loss_name,
         )
         predictions, valid = _align_to_egospeak_timeline(
             probabilities,
@@ -323,6 +338,8 @@ def main() -> None:
                 else None
             ),
             "classes": list(CLASS_NAMES),
+            "loss": loss_name,
+            "output_space": output_space_for_loss(loss_name),
             "prediction_layout": "[target_time, class, future_horizon]",
             "checkpoint": str(args.checkpoint.resolve()),
             "checkpoint_epoch": checkpoint_payload.get("epoch"),