diff --git a/egomimic/eval/eval_latent.py b/egomimic/eval/eval_latent.py index 2d663b3fc..58fa6ba58 100644 --- a/egomimic/eval/eval_latent.py +++ b/egomimic/eval/eval_latent.py @@ -58,6 +58,8 @@ def __init__( compute_pca_umap: bool = True, pca_n_components: int = 50, pca_for_downstream: bool = False, + compute_knn: bool = False, + knn_k: int = 5, emit_combined: bool = True, color_by: str = "embodiment", # "embodiment" or "hash" ): @@ -95,6 +97,8 @@ def __init__( f"color_by must be 'embodiment' or 'hash', got {color_by!r}" ) self.color_by = color_by + self.compute_knn = compute_knn + self.knn_k = knn_k self._layer_keys = {} # layer_name -> list[np.ndarray (B, S, D)] self._row_hashes = [] # one entry per sample (replicated by S at write time) self._row_embodiments = [] @@ -414,6 +418,19 @@ def on_validation_end(self): ): tsne_3d = self._tsne_3d(features_for_reduction) + knn_result = None + if self.compute_knn: + with _timed(f"{layer_name} | KNN-{self.knn_k} ({keys.shape[0]} rows)"): + knn_result = self._knn_accuracy( + features_for_reduction, embs, k=self.knn_k + ) + logger.info( + "[KNN] %s: accuracy=%.4f (+/- %.4f)", + layer_name, + knn_result["accuracy"], + knn_result.get("std", 0.0), + ) + csv_path = os.path.join(out_dir, f"{layer_name}.csv") keys_pt_path = os.path.join(out_dir, f"{layer_name}_keys.pt") with _timed(f"{layer_name} | write_csv ({keys.shape[0]} rows)"): @@ -599,6 +616,38 @@ def _pca(features: np.ndarray, n_components: int): np.float32 ), reducer.explained_variance_ratio_.astype(np.float32) + @staticmethod + def _knn_accuracy(features: np.ndarray, labels: list, k: int = 5) -> dict: + """KNN classification accuracy predicting embodiment from latent features. + Returns dict with accuracy (mean 5-fold CV) and per-fold scores. + Uses cuML GPU KNN when available, sklearn CPU otherwise.""" + n = features.shape[0] + if n < k + 1: + return {"accuracy": 0.0, "per_fold": []} + from sklearn.preprocessing import LabelEncoder + + le = LabelEncoder() + y = le.fit_transform(labels) + if len(np.unique(y)) < 2: + return {"accuracy": 1.0, "per_fold": [1.0]} + X = features.astype(np.float32) + try: + from cuml.neighbors import KNeighborsClassifier + + knn = KNeighborsClassifier(n_neighbors=k, output_type="numpy") + except ImportError: + from sklearn.neighbors import KNeighborsClassifier + + knn = KNeighborsClassifier(n_neighbors=k) + from sklearn.model_selection import cross_val_score + + scores = cross_val_score(knn, X, y, cv=min(5, n), scoring="accuracy") + return { + "accuracy": float(scores.mean()), + "std": float(scores.std()), + "per_fold": scores.tolist(), + } + @staticmethod def _write_csv( path, diff --git a/egomimic/hydra_configs/evaluator/eval_latent.yaml b/egomimic/hydra_configs/evaluator/eval_latent.yaml index f3b464390..8dc48bf80 100644 --- a/egomimic/hydra_configs/evaluator/eval_latent.yaml +++ b/egomimic/hydra_configs/evaluator/eval_latent.yaml @@ -33,3 +33,9 @@ emit_combined: true # good for cotrain_pi_latent_random. Override to "hash" for the pairs config # so each of the 4 specific episodes gets its own color. color_by: embodiment + +# KNN accuracy (embodiment classification from latent features). +# Uses cuML GPU KNN when available, sklearn CPU otherwise. +# Reports 5-fold cross-validation accuracy per layer slice. +compute_knn: false +knn_k: 5