diff --git a/stable_pretraining/callbacks/knn.py b/stable_pretraining/callbacks/knn.py index ade0015c6..ad4e14316 100644 --- a/stable_pretraining/callbacks/knn.py +++ b/stable_pretraining/callbacks/knn.py @@ -47,6 +47,8 @@ class OnlineKNN(Callback): to compute all distances at once. Default is -1. distance_metric: Distance metric for finding nearest neighbors. Options are 'euclidean', 'squared_euclidean', 'cosine', 'manhattan'. Default is 'euclidean'. + multi_label: Whether the dataset uses multi-label classification. Default is False. + num_classes: Number of total classes. Required if multi_label is True. Raises: ValueError: If k <= 0, temperature <= 0, or chunk_size is invalid. @@ -73,6 +75,8 @@ def __init__( distance_metric: Literal[ "euclidean", "squared_euclidean", "cosine", "manhattan" ] = "euclidean", + multi_label: bool = False, + num_classes: Optional[int] = None, ) -> None: super().__init__() @@ -83,6 +87,9 @@ def __init__( if chunk_size == 0 or chunk_size < -1: raise ValueError(f"chunk_size must be positive or -1, got {chunk_size}") + if multi_label and num_classes is None: + raise ValueError("You must provide `num_classes` when `multi_label=True`.") + if input_dim is not None and isinstance(input_dim, (list, tuple)): input_dim = int(np.prod(input_dim)) @@ -97,10 +104,20 @@ def __init__( self.chunk_size = chunk_size self.distance_metric = distance_metric self.metrics = metrics + self.multi_label = multi_label + self.num_classes = num_classes self._input_queue = None self._target_queue = None + for metric in self.metrics.values(): + if metric.__class__.__name__ in ["MultilabelAUROC", "MulticlassAUROC"]: + logging.warning( + f"[{self.name}] AUROC metric detected. torchmetrics expects target labels to be Long integers. " + "If you are passing float targets, you can cast them to long in your dataset/forward pass. " + ) + break + @property def state_key(self) -> str: """Unique identifier for this callback's state during checkpointing.""" @@ -195,11 +212,6 @@ def _compute_knn_predictions( ) -> Optional[Tensor]: """Compute KNN predictions.""" batch_size = features.size(0) - num_classes = int(cached_labels.max().item()) + 1 - - predictions = torch.zeros( - batch_size, num_classes, device=features.device, dtype=torch.float32 - ) if cached_features.device != features.device: cached_features = cached_features.to(features.device) @@ -223,13 +235,21 @@ def _compute_knn_predictions( dist_weight = 1 / dist_weight.add_(self.temperature) - labels_1d = ( - cached_labels.squeeze(-1) if cached_labels.dim() > 1 else cached_labels - ) - selected_labels = labels_1d[sim_indices].long() - one_hot_labels = F.one_hot(selected_labels, num_classes=num_classes) + if self.multi_label: + selected_labels = cached_labels[sim_indices].float() + weighted_votes = dist_weight.unsqueeze(-1) * selected_labels + total_votes = weighted_votes.sum(dim=0) + sum_of_weights = dist_weight.sum(dim=0).unsqueeze(-1) + predictions = total_votes / sum_of_weights + else: + num_classes = int(cached_labels.max().item()) + 1 + labels_1d = ( + cached_labels.squeeze(-1) if cached_labels.dim() > 1 else cached_labels + ) + selected_labels = labels_1d[sim_indices].long() + one_hot_labels = F.one_hot(selected_labels, num_classes=num_classes) + predictions = (dist_weight.unsqueeze(-1) * one_hot_labels).sum(0) - predictions = (dist_weight.unsqueeze(-1) * one_hot_labels).sum(0) return predictions def _log_metrics( diff --git a/stable_pretraining/callbacks/probe.py b/stable_pretraining/callbacks/probe.py index 372ee4099..fed8b5f5f 100644 --- a/stable_pretraining/callbacks/probe.py +++ b/stable_pretraining/callbacks/probe.py @@ -108,6 +108,15 @@ def __init__( self.metrics = metrics logging.info(f"{self.name}: We are wrapping up your `forward`!") self.wrap_forward(pl_module=module) + for metric in self.metrics.values(): + if metric.__class__.__name__ in ["MultilabelAUROC", "MulticlassAUROC"]: + logging.warning( + f"[{self.name}] AUROC metric detected. torchmetrics expects target labels to be Long integers. " + "If you are passing float targets, you can cast them to long in your dataset/forward pass. " + "If your loss function requires floats, you can handle it there (e.g., " + "`loss=lambda p, t: torch.nn.functional.binary_cross_entropy_with_logits(p, t.float())`)." + ) + break def configure_model(self, pl_module: LightningModule) -> torch.nn.Module: """Initialize the probe module from configuration.""" diff --git a/stable_pretraining/tests/unit/test_probing.py b/stable_pretraining/tests/unit/test_probing.py index 28de3a5d3..36d84cac2 100644 --- a/stable_pretraining/tests/unit/test_probing.py +++ b/stable_pretraining/tests/unit/test_probing.py @@ -249,3 +249,77 @@ def test_online_probe_lifecycle_methods(self): assert len(module.callbacks_metrics) == 1 assert len(module.callbacks_modules) == 1 module.configure_optimizers() + + @patch("stable_pretraining.callbacks.knn.logging.warning") + def test_online_knn_multilabel_initialization(self, mock_warning): + """Test multi-label initialization, required args, and AUROC warnings.""" + from stable_pretraining.callbacks import OnlineKNN + import torchmetrics + + # Test missing num_classes raises ValueError + with pytest.raises(ValueError, match="You must provide `num_classes`"): + OnlineKNN( + name="knn_probe", + input="embedding", + target="label", + queue_length=100, + metrics={}, + multi_label=True, + ) + + # Test AUROC warning is triggered properly + metrics = {"auroc": torchmetrics.classification.MultilabelAUROC(num_labels=3)} + knn = OnlineKNN( + name="knn_probe", + input="embedding", + target="label", + queue_length=100, + metrics=metrics, + multi_label=True, + num_classes=3, + ) + + assert knn.multi_label is True + assert knn.num_classes == 3 + mock_warning.assert_called_once() + assert "AUROC metric detected" in mock_warning.call_args[0][0] + + def test_online_knn_multilabel_math(self): + """Test the multi-label KNN distance weighting and voting math.""" + from stable_pretraining.callbacks import OnlineKNN + + knn = OnlineKNN( + name="knn_probe", + input="embedding", + target="label", + queue_length=10, + metrics={}, + k=2, + temperature=1.0, # dist_weight = 1 / (dist + 1) + multi_label=True, + num_classes=3, + ) + + # Setup mock data + query = torch.tensor([[0.0, 0.0]]) + cached_features = torch.tensor( + [ + [0.0, 0.0], # dist=0 -> weight=1.0 + [2.0, 0.0], # dist=2 -> weight=0.3333 + [10.0, 10.0], # Ignored because k=2 + ] + ) + + cached_labels = torch.tensor( + [[1.0, 0.0, 1.0], [1.0, 1.0, 0.0], [0.0, 0.0, 0.0]] + ) + + preds = knn._compute_knn_predictions(query, cached_features, cached_labels) + + # Expected votes: + # Class 0: (1.0 * 1) + (0.333 * 1) = 1.333 / 1.333 = 1.0 + # Class 1: (1.0 * 0) + (0.333 * 1) = 0.333 / 1.333 = 0.25 + # Class 2: (1.0 * 1) + (0.333 * 0) = 1.000 / 1.333 = 0.75 + expected = torch.tensor([[1.00, 0.25, 0.75]]) + + assert torch.allclose(preds, expected, atol=1e-3)