XminorAbi/SelfSupervisedLearning

1"""2cluster.py — K-Means Clustering + Evaluation3=============================================4Team 14: Abhinandan Chakraborty & Kamal Kishor Dhakad5Project: Learning Visual Concepts Without Labels6 7FIXES APPLIED:8--------------9FIX 2 — Representation mismatch:10  Embeddings are now L2-normalised h (from model.encode()).11  K-Means is run on the normalised hypersphere — consistent with training.12 13FIX 3 — Distance metric consistency:14  K-Means uses Euclidean distance. On the unit hypersphere:15    ||u - v||² = 2 - 2·cos(u,v)16  So Euclidean distance on normalised vectors is monotone with cosine distance.17  This makes K-Means on normalised embeddings equivalent to spherical K-Means.18  No more metric mismatch.19 20FIX 6 — Dimensionality reduction:21  Reduced PCA from 512→50 to 512→128.22  At 50 dims, variance explained ≈ 82% (losing 18% of structure).23  At 128 dims, variance explained ≈ 95%+ (preserving fine-grained cluster structure).24  K-Means on normalised embeddings is fast enough at 128 dims.25  If runtime is an issue, reduce to 64 (not 50).26"""27 28import torch29import numpy as np30import matplotlib.pyplot as plt31from pathlib import Path32from sklearn.cluster import KMeans33from sklearn.decomposition import PCA34from sklearn.manifold import TSNE35from sklearn.preprocessing import normalize36from sklearn.metrics import normalized_mutual_info_score, adjusted_rand_score37from scipy.optimize import linear_sum_assignment38 39 40def hungarian_accuracy(true_labels, pred_labels) -> float:41    """Clustering accuracy with Hungarian matching."""42    n  = max(true_labels.max(), pred_labels.max()) + 143    cm = np.zeros((n, n), dtype=np.int64)44    for t, p in zip(true_labels, pred_labels):45        cm[p, t] += 146    r, c = linear_sum_assignment(-cm)47    return cm[r, c].sum() / len(true_labels)48 49 50def cluster_embeddings(51    embeddings_path: str  = "embeddings/embeddings.pt",52    labels_path    : str  = "embeddings/labels.pt",53    n_clusters     : int  = 10,54    n_init         : int  = 20,55    pca_dims       : int  = 128,    # FIX 6: 128 instead of 5056    results_dir    : str  = "results",57    class_names    : list = None,58) -> dict:59    """60    Full clustering pipeline with all fixes applied:61      Load L2-normalised embeddings62      → Re-normalise (safety) → PCA(128) → K-Means → Hungarian → t-SNE → plots63    """64    from dataset import CIFAR10_CLASSES65    Path(results_dir).mkdir(exist_ok=True)66 67    if class_names is None:68        class_names = CIFAR10_CLASSES69 70    print("[Clustering] Loading embeddings...")71    embeddings = torch.load(embeddings_path).numpy()   # (N, 512) — L2-normalised72    labels     = torch.load(labels_path).numpy()       # (N,)73    N          = len(labels)74    print(f"  {N:,} embeddings | dim={embeddings.shape[1]}")75 76    # FIX 2+3: Re-normalise to guarantee unit norm (safety check)77    embeddings = normalize(embeddings, norm="l2")78    mean_norm  = np.linalg.norm(embeddings, axis=1).mean()79    print(f"  Mean L2 norm after normalisation: {mean_norm:.4f} (should be 1.0)")80 81    # FIX 6: PCA 512 → 128 (not 50)82    print(f"[Clustering] PCA: {embeddings.shape[1]} → {pca_dims} dims...")83    pca    = PCA(n_components=pca_dims, random_state=42)84    emb_pca = pca.fit_transform(embeddings)85    var    = pca.explained_variance_ratio_.sum()86    print(f"  Variance explained by {pca_dims} PCs: {var:.1%}")87 88    # FIX 3: Re-normalise after PCA (PCA destroys unit norm)89    # This keeps K-Means consistent with cosine distance90    emb_pca = normalize(emb_pca, norm="l2")91 92    # K-Means on normalised PCA features93    print(f"[Clustering] K-Means K={n_clusters}, n_init={n_init}...")94    kmeans = KMeans(95        n_clusters=n_clusters, init="k-means++",96        n_init=n_init, max_iter=500, random_state=42,97    )98    cluster_assignments = kmeans.fit_predict(emb_pca)99 100    # Evaluate101    acc = hungarian_accuracy(labels, cluster_assignments)102    nmi = normalized_mutual_info_score(labels, cluster_assignments)103    ari = adjusted_rand_score(labels, cluster_assignments)104 105    print(f"\n{'='*50}")106    print(f"  Clustering Accuracy (Hungarian) : {acc*100:.1f}%")107    print(f"  NMI                             : {nmi:.4f}")108    print(f"  ARI                             : {ari:.4f}")109    print(f"{'='*50}")110 111    # t-SNE on normalised PCA features112    print("\n[Clustering] t-SNE (1-3 mins)...")113    perp  = min(30, N // 4)114    tsne  = TSNE(n_components=2, perplexity=perp, learning_rate="auto",115                 init="pca", random_state=42, max_iter=1000)116    emb_2d = tsne.fit_transform(emb_pca)117 118    _plot_embeddings(emb_2d, labels, class_names,119                     "t-SNE — ground truth labels",120                     f"{results_dir}/tsne_true_labels.png")121    _plot_embeddings(emb_2d, cluster_assignments,122                     [f"Cluster {i}" for i in range(n_clusters)],123                     f"t-SNE — K-Means clusters (K={n_clusters})",124                     f"{results_dir}/tsne_clusters.png")125    _plot_metrics(acc, nmi, ari, f"{results_dir}/metrics.png")126 127    print(f"[Clustering] Plots saved to {results_dir}/")128 129    return {130        "accuracy": acc, "nmi": nmi, "ari": ari,131        "cluster_assignments": cluster_assignments,132        "embeddings_2d": emb_2d,133        "kmeans": kmeans, "pca": pca,134    }135 136 137def _plot_embeddings(coords, labels, class_names, title, save_path):138    colors = plt.cm.tab10(np.linspace(0, 1, len(class_names)))139    fig, ax = plt.subplots(figsize=(9, 7))140    for i, (name, color) in enumerate(zip(class_names, colors)):141        mask = labels == i142        ax.scatter(coords[mask, 0], coords[mask, 1],143                   c=[color], label=name, s=8, alpha=0.6, edgecolors='none')144    ax.set_title(title, fontsize=12, pad=10)145    ax.set_xticks([]); ax.set_yticks([])146    ax.legend(loc="best", fontsize=8, markerscale=2)147    plt.tight_layout()148    plt.savefig(save_path, dpi=150, bbox_inches='tight')149    plt.close()150    print(f"  Saved: {save_path}")151 152 153def _plot_metrics(acc, nmi, ari, save_path):154    fig, ax = plt.subplots(figsize=(5, 4))155    vals   = [acc, nmi, ari]156    names  = ["Accuracy", "NMI", "ARI"]157    colors = ["#534AB7", "#1D9E75", "#D85A30"]158    bars   = ax.bar(names, vals, color=colors, width=0.5, zorder=3)159    ax.set_ylim(0, 1.1)160    ax.set_ylabel("Score")161    ax.set_title("Clustering metrics — CIFAR-10", fontsize=11)162    ax.yaxis.grid(True, linestyle='--', alpha=0.4, zorder=0)163    for bar, val in zip(bars, vals):164        ax.text(bar.get_x() + bar.get_width()/2,165                bar.get_height() + 0.02, f"{val:.3f}", ha='center', fontsize=11)166    plt.tight_layout()167    plt.savefig(save_path, dpi=150, bbox_inches='tight')168    plt.close()169    print(f"  Saved: {save_path}")170 171 172if __name__ == "__main__":173    cluster_embeddings()174