diff --git a/README.md b/README.md index d4d0ffc..e21773e 100644 --- a/README.md +++ b/README.md @@ -954,6 +954,21 @@ without normalizing the target. The previous L2-normalized objective belongs to MiniROAD and is available only as the explicit `--loss miniroad-normalized` ablation. +Two label-structure ablations are also available without changing the saved +three-column target files: + +- `--loss independent-bce` keeps a `10 x 3` head, trains the three labels with + independent sigmoid/BCE, and evaluates the sigmoid marginals directly; +- `--loss four-state-ce` changes the head to `10 x 4` and maps each target to + `silence`, `wearer_only`, `others_only`, or `overlap`. At evaluation time it + returns the original three marginal scores as `P(silence)`, + `P(wearer_only)+P(overlap)`, and `P(others_only)+P(overlap)`. + +The four-state marginals are intentionally not renormalized. Keep features, +split, fusion, optimizer, schedule, and seeds fixed when comparing these losses. +For checkpoint selection, reserve one official training session such as +Session 12 with `--dev-sessions 12`; Sessions 1--3 remain test-only. + Recommended experimental procedure: 1. Freeze Qwen and extract all training-session states offline. diff --git a/easycom_head.py b/easycom_head.py index 0dc842b..ba2b2b9 100644 --- a/easycom_head.py +++ b/easycom_head.py @@ -30,9 +30,22 @@ CLASS_NAMES = ( "camera_wearer_speaking", "other_people_speaking", ) +FOUR_STATE_NAMES = ( + "silence", + "wearer_only", + "others_only", + "overlap", +) STEP_SECONDS = 0.2 NUM_HORIZONS = 10 NUM_CLASSES = 3 +NUM_FOUR_STATES = len(FOUR_STATE_NAMES) +LOSS_NAMES = ( + "lstr-mce", + "miniroad-normalized", + "independent-bce", + "four-state-ce", +) DATASET_NAMES = ("easycom", "ego4d") LAYER_FUSION_MODES = ( "single", @@ -44,7 +57,7 @@ DatasetName = Literal["easycom", "ego4d"] class TurnPredictionHead(nn.Module): - """Map one Omni state to EgoSpeak-compatible ``[10, 3]`` logits.""" + """Map one Omni state to configurable ``[horizon, output]`` logits.""" def __init__( self, @@ -150,18 +163,76 @@ class TurnPredictionHead(nn.Module): return self.layer_logits.softmax(dim=0) -def _validate_loss_inputs(logits: torch.Tensor, targets: torch.Tensor) -> None: +def _validate_loss_inputs( + logits: torch.Tensor, + targets: torch.Tensor, + *, + num_output_classes: int = NUM_CLASSES, +) -> None: if logits.shape != targets.shape: raise ValueError( f"logits/targets shape mismatch: {tuple(logits.shape)} vs " f"{tuple(targets.shape)}" ) - if logits.shape[-1] != NUM_CLASSES: + if logits.shape[-1] != num_output_classes: raise ValueError( - f"Expected {NUM_CLASSES} classes, got {logits.shape[-1]}" + f"Expected {num_output_classes} classes, got {logits.shape[-1]}" ) +def _validated_binary_activity_targets(targets: torch.Tensor) -> torch.Tensor: + """Validate the four legal EasyCom/EgoSpeak three-label activity rows.""" + + if targets.ndim < 1 or targets.shape[-1] != NUM_CLASSES: + raise ValueError( + f"Expected activity targets [...,{NUM_CLASSES}], got " + f"{tuple(targets.shape)}" + ) + if not torch.isfinite(targets).all(): + raise ValueError("Activity targets contain NaN or Inf") + rounded = targets.round() + if not torch.allclose(targets, rounded, atol=1e-6, rtol=0.0): + raise ValueError("Activity targets must contain only binary 0/1 values") + expected_background = (rounded[..., 1:].sum(dim=-1) == 0).to(rounded.dtype) + if not torch.equal(rounded[..., 0], expected_background): + raise ValueError( + "Activity targets must be one of [1,0,0], [0,1,0], " + "[0,0,1], or [0,1,1]" + ) + return rounded + + +def activity_targets_to_four_states(targets: torch.Tensor) -> torch.LongTensor: + """Map three-label activity targets to silence/wearer/others/overlap.""" + + targets = _validated_binary_activity_targets(targets) + return (targets[..., 1] + 2 * targets[..., 2]).long() + + +def output_classes_for_loss(loss_name: str) -> int: + """Return the number of head logits required by one training objective.""" + + if loss_name not in LOSS_NAMES: + raise ValueError(f"Unknown loss {loss_name!r}; expected one of {LOSS_NAMES}") + return NUM_FOUR_STATES if loss_name == "four-state-ce" else NUM_CLASSES + + +def output_space_for_loss(loss_name: str) -> str: + """Return a stable checkpoint description of the head output space.""" + + output_classes_for_loss(loss_name) + if loss_name == "four-state-ce": + return "four-state-softmax" + if loss_name == "independent-bce": + return "three-label-sigmoid" + return "three-class-softmax" + + +def head_output_names_for_loss(loss_name: str) -> tuple[str, ...]: + output_classes_for_loss(loss_name) + return FOUR_STATE_NAMES if loss_name == "four-state-ce" else CLASS_NAMES + + def egospeak_multihorizon_loss( logits: torch.Tensor, targets: torch.Tensor, @@ -211,6 +282,99 @@ def miniroad_multihorizon_loss( raise ValueError(f"Unknown reduction: {reduction}") +def independent_bce_multihorizon_loss( + logits: torch.Tensor, + targets: torch.Tensor, + *, + reduction: Literal["mean", "sum", "none"] = "mean", +) -> torch.Tensor: + """Independent sigmoid/BCE loss for the three marginal activity labels.""" + + _validate_loss_inputs(logits, targets) + targets = targets.to( + device=logits.device, + dtype=logits.dtype, + ) + element_loss = nn.functional.binary_cross_entropy_with_logits( + logits, + targets, + reduction="none", + ) + if reduction == "mean": + return element_loss.mean() + if reduction == "sum": + return element_loss.sum() + if reduction == "none": + return element_loss + raise ValueError(f"Unknown reduction: {reduction}") + + +def four_state_multihorizon_loss( + logits: torch.Tensor, + targets: torch.Tensor, + *, + reduction: Literal["mean", "sum", "none"] = "mean", +) -> torch.Tensor: + """Four-state CE for silence, wearer-only, others-only, and overlap.""" + + if logits.ndim < 1 or logits.shape[-1] != NUM_FOUR_STATES: + raise ValueError( + f"Expected four-state logits [...,{NUM_FOUR_STATES}], got " + f"{tuple(logits.shape)}" + ) + if logits.shape[:-1] != targets.shape[:-1] or targets.shape[-1] != NUM_CLASSES: + raise ValueError( + "Four-state logits/activity target shape mismatch: " + f"{tuple(logits.shape)} vs {tuple(targets.shape)}" + ) + state_targets = (targets[..., 1] + 2 * targets[..., 2]).long().to( + logits.device + ) + loss = nn.functional.cross_entropy( + logits.reshape(-1, NUM_FOUR_STATES), + state_targets.reshape(-1), + reduction="none", + ).reshape(logits.shape[:-1]) + if reduction == "mean": + return loss.mean() + if reduction == "sum": + return loss.sum() + if reduction == "none": + return loss + raise ValueError(f"Unknown reduction: {reduction}") + + +def marginal_probabilities_from_logits( + logits: torch.Tensor, + loss_name: str, +) -> torch.Tensor: + """Convert any supported head output to three EgoSpeak marginals. + + Four-state marginals deliberately are not renormalized: wearer and others + may both be near one when the overlap state is likely. + """ + + expected_classes = output_classes_for_loss(loss_name) + if logits.ndim < 1 or logits.shape[-1] != expected_classes: + raise ValueError( + f"Loss {loss_name!r} expects {expected_classes} logits, got " + f"{tuple(logits.shape)}" + ) + if loss_name == "independent-bce": + return logits.sigmoid() + probabilities = logits.softmax(dim=-1) + if loss_name != "four-state-ce": + return probabilities + return torch.stack( + ( + probabilities[..., 0], + probabilities[..., 1] + probabilities[..., 3], + probabilities[..., 2] + probabilities[..., 3], + ), + dim=-1, + ) + + def multihorizon_cross_entropy( logits: torch.Tensor, targets: torch.Tensor, @@ -380,6 +544,7 @@ def load_session_data( f"{target_path}" ) targets = torch.from_numpy(np.asarray(target_array, dtype=np.float32)) + _validated_binary_activity_targets(targets) if len(label_indices) and int(label_indices[-1]) >= len(targets): raise ValueError( f"Last label index {int(label_indices[-1])} exceeds target length " @@ -455,6 +620,30 @@ def checkpoint_head_config(payload: dict[str, Any]) -> dict[str, Any]: return config +def checkpoint_loss_name(payload: dict[str, Any]) -> str: + """Read the objective from a checkpoint, defaulting legacy files to MCE.""" + + training_config = payload.get("training_config", {}) + if not isinstance(training_config, dict): + raise ValueError("Checkpoint training_config must be a mapping") + loss_name = str(training_config.get("loss", "lstr-mce")) + expected_classes = output_classes_for_loss(loss_name) + actual_classes = int(checkpoint_head_config(payload)["num_classes"]) + if actual_classes != expected_classes: + raise ValueError( + f"Checkpoint loss {loss_name!r} requires {expected_classes} head " + f"outputs, but head_config has {actual_classes}" + ) + recorded_space = payload.get("output_space") + expected_space = output_space_for_loss(loss_name) + if recorded_space is not None and recorded_space != expected_space: + raise ValueError( + f"Checkpoint output_space is {recorded_space!r}, expected " + f"{expected_space!r} for loss {loss_name!r}" + ) + return loss_name + + def validate_disjoint_sessions( first: Sequence[str], second: Sequence[str], diff --git a/infer_turn_head.py b/infer_turn_head.py index 8f4386d..1af8470 100644 --- a/infer_turn_head.py +++ b/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"), diff --git a/streaming_turn_model.py b/streaming_turn_model.py index f42b604..8c7963a 100644 --- a/streaming_turn_model.py +++ b/streaming_turn_model.py @@ -44,7 +44,9 @@ import torch from easycom_head import ( TurnPredictionHead, checkpoint_head_config, + checkpoint_loss_name, egospeak_multihorizon_loss, + marginal_probabilities_from_logits, miniroad_multihorizon_loss, multihorizon_cross_entropy, ) @@ -2141,10 +2143,11 @@ def _load_head( checkpoint: Path, hidden_size: int, device: torch.device, -) -> TurnPredictionHead: +) -> tuple[TurnPredictionHead, str]: payload = torch.load(checkpoint, map_location="cpu", weights_only=True) if isinstance(payload, dict) and "head" in payload: config = checkpoint_head_config(payload) + loss_name = checkpoint_loss_name(payload) if int(config["hidden_size"]) != hidden_size: raise ValueError( f"Head hidden size {config['hidden_size']} != model {hidden_size}" @@ -2154,8 +2157,9 @@ def _load_head( else: head = TurnPredictionHead(hidden_size=hidden_size) state_dict = payload + loss_name = "lstr-mce" head.load_state_dict(state_dict) - return head.to(device).eval() + return head.to(device).eval(), loss_name def parse_args() -> argparse.Namespace: @@ -2393,11 +2397,19 @@ def main() -> None: "verification": result.verification, } if args.head_checkpoint is not None: - head = _load_head(args.head_checkpoint, _hidden_size(model), model.device) + head, head_loss_name = _load_head( + args.head_checkpoint, + _hidden_size(model), + model.device, + ) with torch.inference_mode(): logits = head(result.states.to(model.device, dtype=model.dtype)).float().cpu() payload["logits"] = logits - payload["probabilities"] = logits.softmax(dim=-1) + payload["probabilities"] = marginal_probabilities_from_logits( + logits, + head_loss_name, + ) + payload["head_loss"] = head_loss_name if args.num_steps is not None and len(result.times) != args.num_steps: raise RuntimeError( diff --git a/streaming_turn_model_directly.py b/streaming_turn_model_directly.py index 23f2ad6..91fcdac 100644 --- a/streaming_turn_model_directly.py +++ b/streaming_turn_model_directly.py @@ -1540,7 +1540,7 @@ def main( "verification": result.verification, } if args.head_checkpoint is not None: - head = core._load_head( + head, head_loss_name = core._load_head( args.head_checkpoint, core._hidden_size(model), model.device ) with torch.inference_mode(): @@ -1548,7 +1548,11 @@ def main( result.states.to(model.device, dtype=model.dtype) ).float().cpu() payload["logits"] = logits - payload["probabilities"] = logits.softmax(dim=-1) + payload["probabilities"] = core.marginal_probabilities_from_logits( + logits, + head_loss_name, + ) + payload["head_loss"] = head_loss_name args.output.parent.mkdir(parents=True, exist_ok=True) torch.save(payload, args.output) diff --git a/test_easycom_head_objectives.py b/test_easycom_head_objectives.py new file mode 100644 index 0000000..523410c --- /dev/null +++ b/test_easycom_head_objectives.py @@ -0,0 +1,109 @@ +"""Focused tests for EasyCom loss/output-space ablations.""" + +from __future__ import annotations + +import math +import unittest + +import torch + +from easycom_head import ( + TurnPredictionHead, + activity_targets_to_four_states, + checkpoint_loss_name, + four_state_multihorizon_loss, + independent_bce_multihorizon_loss, + marginal_probabilities_from_logits, + output_classes_for_loss, +) + + +class ObjectiveTests(unittest.TestCase): + def setUp(self) -> None: + self.targets = torch.tensor( + [ + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + [0.0, 1.0, 1.0], + ] + ) + + def test_activity_target_mapping(self) -> None: + actual = activity_targets_to_four_states(self.targets) + self.assertTrue(torch.equal(actual, torch.tensor([0, 1, 2, 3]))) + + def test_invalid_activity_target_is_rejected(self) -> None: + with self.assertRaisesRegex(ValueError, "must be one of"): + activity_targets_to_four_states(torch.tensor([[0.0, 0.0, 0.0]])) + + def test_independent_bce_uses_standard_mean(self) -> None: + logits = torch.zeros_like(self.targets) + loss = independent_bce_multihorizon_loss(logits, self.targets) + self.assertAlmostEqual(float(loss), math.log(2.0), places=6) + + def test_four_state_loss_accepts_three_label_targets(self) -> None: + logits = torch.zeros(4, 4) + loss = four_state_multihorizon_loss(logits, self.targets) + self.assertAlmostEqual(float(loss), math.log(4.0), places=6) + + def test_four_state_probabilities_are_marginalized(self) -> None: + state_probabilities = torch.tensor([[0.1, 0.2, 0.3, 0.4]]) + logits = state_probabilities.log() + actual = marginal_probabilities_from_logits(logits, "four-state-ce") + expected = torch.tensor([[0.1, 0.6, 0.7]]) + self.assertTrue(torch.allclose(actual, expected, atol=1e-6)) + self.assertGreater(float(actual.sum()), 1.0) + + def test_bce_probabilities_use_sigmoid(self) -> None: + actual = marginal_probabilities_from_logits( + torch.zeros(2, 3), + "independent-bce", + ) + self.assertTrue(torch.equal(actual, torch.full((2, 3), 0.5))) + + def test_four_state_head_shape(self) -> None: + head = TurnPredictionHead( + hidden_size=8, + num_horizons=2, + num_classes=output_classes_for_loss("four-state-ce"), + dropout=0.0, + ) + self.assertEqual(tuple(head(torch.randn(3, 8)).shape), (3, 2, 4)) + + def test_all_new_objectives_backpropagate_to_marginals(self) -> None: + states = torch.randn(4, 8) + future_targets = self.targets[:, None, :].repeat(1, 2, 1) + for loss_name, loss_function in ( + ("independent-bce", independent_bce_multihorizon_loss), + ("four-state-ce", four_state_multihorizon_loss), + ): + with self.subTest(loss=loss_name): + head = TurnPredictionHead( + hidden_size=8, + num_horizons=2, + num_classes=output_classes_for_loss(loss_name), + dropout=0.0, + ) + logits = head(states) + loss = loss_function(logits, future_targets) + loss.backward() + scores = marginal_probabilities_from_logits(logits, loss_name) + self.assertEqual(tuple(scores.shape), (4, 2, 3)) + self.assertTrue(torch.isfinite(scores).all()) + self.assertIsNotNone(head.net[-1].weight.grad) + + def test_legacy_checkpoint_defaults_to_lstr(self) -> None: + self.assertEqual(checkpoint_loss_name({}), "lstr-mce") + + def test_checkpoint_rejects_loss_output_mismatch(self) -> None: + payload = { + "training_config": {"loss": "four-state-ce"}, + "head_config": {"num_classes": 3}, + } + with self.assertRaisesRegex(ValueError, "requires 4"): + checkpoint_loss_name(payload) + + +if __name__ == "__main__": + unittest.main() diff --git a/train_turn_head.py b/train_turn_head.py index cecebaf..c03cbf4 100644 --- a/train_turn_head.py +++ b/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