Download scripts/plot_cosne.py from AIGeeksGroup/HMCLIP: direct link, hf CLI and curl.
- Browser
- Download file 23.3 kB
-
https://huggingface.co/AIGeeksGroup/HMCLIP/resolve/main/scripts/plot_cosne.py
- Command line
-
hf download hf://AIGeeksGroup/HMCLIP/scripts/plot_cosne.py
-
curl -L -o plot_cosne.py https://huggingface.co/AIGeeksGroup/HMCLIP/resolve/main/scripts/plot_cosne.py
23.3 kB
| #!/usr/bin/env python3 | |
| """ | |
| CO-SNE visualization for RadIR/RadIR *hyperbolic* embeddings. | |
| Unlike scripts/plot_tsne.py (log_map0 → Euclidean t-SNE), this keeps points in | |
| hyperbolic geometry: Lorentz space components are mapped to the Poincaré disk, | |
| then CO-SNE (CVPR 2022) reduces them to 2D while preserving: | |
| - local hyperbolic neighborhoods | |
| - each point's geodesic distance to the origin (hierarchy / radius) | |
| Reference: Guo et al., "CO-SNE: Dimensionality Reduction and Visualization for | |
| Hyperbolic Data", CVPR 2022. https://github.com/yunhuiguo/CO-SNE | |
| This file ships a self-contained CO-SNE core (modern PyTorch) so we do not | |
| depend on the official repo's old sklearn / hyptorch pins. | |
| Usage (from RadIR/RadIR): | |
| python scripts/plot_cosne.py \\ | |
| --load_feats outputs/tsne_feats.npz \\ | |
| --out outputs/cosne_radir.png | |
| # Smaller / faster smoke run: | |
| python scripts/plot_cosne.py --load_feats outputs/tsne_feats.npz \\ | |
| --max_samples 200 --n_iter 500 --out outputs/cosne_radir_test.png | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import os | |
| import sys | |
| import matplotlib.pyplot as plt | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| from torch.utils.data import DataLoader | |
| from transformers import AutoModel, AutoTokenizer | |
| ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) | |
| RAD_IR_ROOT = os.path.join(ROOT, "Rad_IR") | |
| for p in (ROOT, RAD_IR_ROOT): | |
| if p not in sys.path: | |
| sys.path.insert(0, p) | |
| from radir import lorentz as L | |
| from radir.RadIR import RADIR, l2norm | |
| from radir.data.hmclip_mimic_box import HyCoClipMimicBoxDataset_JPG, collate_hmclip_box | |
| from transformer_maskgit.transformer_maskgit.ctvit import CTViT | |
| EPS = 1e-8 | |
| BALL_EPS = 1e-5 | |
| # --------------------------------------------------------------------------- | |
| # Lorentz (RadIR) → Poincaré ball (CO-SNE native space) | |
| # --------------------------------------------------------------------------- | |
| def lorentz_space_to_poincare(z_space: torch.Tensor, curv: float) -> torch.Tensor: | |
| """ | |
| Map Lorentz *space* components (RadIR/MERU) to the Poincaré ball of | |
| curvature -curv, then rescale into the unit disk for visualization. | |
| z_space: [N, D] | |
| returns: [N, D] with ||u|| < 1 | |
| """ | |
| c = float(max(curv, EPS)) | |
| sqrt_c = c**0.5 | |
| # time component: t = sqrt(1/c + ||z||^2) | |
| t = torch.sqrt(1.0 / c + z_space.pow(2).sum(dim=-1, keepdim=True)) | |
| # Poincaré coords in ball of radius 1/sqrt(c): | |
| # u = z / (t + 1/sqrt(c)) | |
| u = z_space / (t + 1.0 / sqrt_c) | |
| # rescale to unit disk: û = sqrt(c) * u | |
| u = sqrt_c * u | |
| return project_to_ball(u) | |
| def project_to_ball(u: torch.Tensor, eps: float = BALL_EPS) -> torch.Tensor: | |
| """Project points into the open unit disk.""" | |
| norm = u.norm(dim=-1, keepdim=True).clamp_min(EPS) | |
| max_norm = 1.0 - eps | |
| scale = torch.clamp(max_norm / norm, max=1.0) | |
| return u * scale | |
| def poincare_dist(x: torch.Tensor, y: torch.Tensor, eps: float = EPS) -> torch.Tensor: | |
| """ | |
| Pairwise geodesic distance on the unit Poincaré disk (curvature -1). | |
| x: [N, D], y: [M, D] → [N, M] | |
| """ | |
| x2 = x.pow(2).sum(dim=-1, keepdim=True) # [N,1] | |
| y2 = y.pow(2).sum(dim=-1, keepdim=True).T # [1,M] | |
| xy = x @ y.T | |
| num = (x2 + y2 - 2.0 * xy).clamp_min(0.0) | |
| den = ((1.0 - x2).clamp_min(eps) * (1.0 - y2).clamp_min(eps)) | |
| z = 1.0 + 2.0 * num / den | |
| return torch.acosh(z.clamp_min(1.0 + eps)) | |
| def poincare_dist_to_origin(x: torch.Tensor, eps: float = EPS) -> torch.Tensor: | |
| """Geodesic distance from each point to the origin on the unit disk.""" | |
| r2 = x.pow(2).sum(dim=-1).clamp(max=1.0 - eps) | |
| return torch.atanh(torch.sqrt(r2).clamp(max=1.0 - eps)) | |
| # --------------------------------------------------------------------------- | |
| # CO-SNE core (self-contained) | |
| # --------------------------------------------------------------------------- | |
| def _binary_search_perplexity( | |
| dist_row: np.ndarray, | |
| perplexity: float, | |
| tol: float = 1e-5, | |
| max_iter: int = 50, | |
| ) -> np.ndarray: | |
| """Given distances to other points, find Gaussian bandwidth for target perplexity.""" | |
| n = dist_row.shape[0] | |
| beta = 1.0 # = 1/(2σ^2) | |
| betamin, betamax = -np.inf, np.inf | |
| log_perp = np.log(perplexity) | |
| for _ in range(max_iter): | |
| p = np.exp(-dist_row * beta) | |
| sum_p = np.maximum(p.sum(), EPS) | |
| p = p / sum_p | |
| entropy = -np.sum(p * np.log(np.maximum(p, EPS))) | |
| diff = entropy - log_perp | |
| if abs(diff) < tol: | |
| break | |
| if diff > 0: | |
| betamin = beta | |
| beta = beta * 2.0 if np.isinf(betamax) else (beta + betamax) / 2.0 | |
| else: | |
| betamax = beta | |
| beta = beta / 2.0 if np.isinf(betamin) else (beta + betamin) / 2.0 | |
| p = np.exp(-dist_row * beta) | |
| p = p / np.maximum(p.sum(), EPS) | |
| return p.astype(np.float64) | |
| def joint_probabilities_from_distances( | |
| dist: np.ndarray, perplexity: float | |
| ) -> np.ndarray: | |
| """ | |
| Build symmetric joint probabilities P from a hyperbolic distance matrix, | |
| matching the t-SNE / CO-SNE high-dim affinity construction. | |
| """ | |
| n = dist.shape[0] | |
| P = np.zeros((n, n), dtype=np.float64) | |
| for i in range(n): | |
| d = dist[i].copy() | |
| d[i] = np.inf | |
| P[i] = _binary_search_perplexity(d, perplexity) | |
| P = P + P.T | |
| P = np.maximum(P / np.maximum(P.sum(), EPS), EPS) | |
| np.fill_diagonal(P, 0.0) | |
| return P | |
| class COSNE: | |
| """ | |
| CO-SNE: hyperbolic t-SNE with origin-distance preservation. | |
| High-dim affinities from hyperbolic geodesic distances. | |
| Low-dim affinities: hyperbolic Cauchy Q_ij ∝ (1 + (d_ij/γ)^2)^-1 | |
| Extra loss: || d_H(x_i,0) - d_H(y_i,0) ||^2 (hierarchy / radius) | |
| """ | |
| def __init__( | |
| self, | |
| n_components: int = 2, | |
| perplexity: float = 30.0, | |
| n_iter: int = 1000, | |
| learning_rate: float = 1.0, | |
| learning_rate_h: float = 0.1, | |
| early_exaggeration: float = 4.0, | |
| student_t_gamma: float = 0.1, | |
| exaggeration_iters: int = 250, | |
| seed: int = 42, | |
| device: str = "cpu", | |
| verbose: bool = True, | |
| ): | |
| self.n_components = n_components | |
| self.perplexity = perplexity | |
| self.n_iter = n_iter | |
| self.learning_rate = learning_rate | |
| self.learning_rate_h = learning_rate_h | |
| self.early_exaggeration = early_exaggeration | |
| self.student_t_gamma = student_t_gamma | |
| self.exaggeration_iters = exaggeration_iters | |
| self.seed = seed | |
| self.device = device | |
| self.verbose = verbose | |
| def fit_transform(self, X_poincare: np.ndarray) -> np.ndarray: | |
| """ | |
| X_poincare: [N, D] points already on the unit Poincaré disk. | |
| Returns Y: [N, 2] on the unit disk. | |
| """ | |
| X = torch.as_tensor(X_poincare, dtype=torch.float64) | |
| X = project_to_ball(X) | |
| n, d = X.shape | |
| device = torch.device(self.device) | |
| if self.verbose: | |
| print(f"[CO-SNE] n={n}, dim={d}, perplexity={self.perplexity}") | |
| # --- high-dim hyperbolic distances & P --- | |
| with torch.no_grad(): | |
| dist_h = poincare_dist(X, X).cpu().numpy() | |
| np.fill_diagonal(dist_h, 0.0) | |
| max_perp = max(2.0, (n - 1) / 3.0) | |
| perp = min(self.perplexity, max_perp) | |
| P = joint_probabilities_from_distances(dist_h, perp) | |
| P = torch.as_tensor(P, dtype=torch.float64, device=device) | |
| # origin distances in high-dim (hierarchy target) | |
| with torch.no_grad(): | |
| d0_high = poincare_dist_to_origin(X).to(device) | |
| # --- init low-dim embedding from first Poincaré coords + radius match --- | |
| rng = np.random.default_rng(self.seed) | |
| if d >= 2: | |
| Y0 = X[:, :2].cpu().numpy() + rng.normal(0, 1e-4, size=(n, 2)) | |
| else: | |
| Y0 = rng.normal(0, 1e-3, size=(n, self.n_components)) | |
| r_high = np.tanh(d0_high.cpu().numpy()) # (0,1) proxy for hierarchy radius | |
| r_y = np.linalg.norm(Y0, axis=1) + EPS | |
| Y0 = Y0 * (r_high / r_y)[:, None] | |
| Y = torch.as_tensor(Y0, dtype=torch.float64, device=device) | |
| Y = project_to_ball(Y) | |
| Y = torch.nn.Parameter(Y) | |
| opt = torch.optim.Adam([Y], lr=self.learning_rate) | |
| # after early exaggeration, decay LR a bit for radius stability | |
| sched = torch.optim.lr_scheduler.StepLR(opt, step_size=max(self.exaggeration_iters, 1), gamma=0.5) | |
| gamma = float(self.student_t_gamma) | |
| for it in range(self.n_iter): | |
| opt.zero_grad() | |
| Y_proj = project_to_ball(Y) | |
| # hyperbolic distances in 2D | |
| dist_y = poincare_dist(Y_proj, Y_proj) | |
| # hyperbolic Cauchy affinities (CO-SNE / HTSNE style) | |
| q_num = 1.0 / (1.0 + (dist_y / (gamma**2 + EPS)).pow(2)) | |
| q_num = q_num.fill_diagonal_(0.0) | |
| Q = q_num / q_num.sum().clamp_min(EPS) | |
| Q = Q.clamp_min(EPS) | |
| # early exaggeration on P | |
| if it < self.exaggeration_iters: | |
| P_use = P * self.early_exaggeration | |
| P_use = P_use / P_use.sum() | |
| else: | |
| P_use = P | |
| kl = (P_use * (torch.log(P_use.clamp_min(EPS)) - torch.log(Q))).sum() | |
| # hierarchy loss: match Poincaré *radii* (more stable than raw d_H | |
| # when high-dim radii are small, as with RadIR's near-constant image norms) | |
| r_low = Y_proj.norm(dim=-1).clamp(max=1.0 - BALL_EPS) | |
| r_high = torch.tanh(d0_high).clamp(max=1.0 - BALL_EPS) | |
| h_loss = F.mse_loss(r_low, r_high) | |
| loss = kl + self.learning_rate_h * h_loss | |
| loss.backward() | |
| # clip grads — Poincaré optimization is sensitive near the boundary | |
| torch.nn.utils.clip_grad_norm_([Y], max_norm=1.0) | |
| opt.step() | |
| with torch.no_grad(): | |
| Y.data.copy_(project_to_ball(Y.data)) | |
| sched.step() | |
| if self.verbose and (it % 100 == 0 or it == self.n_iter - 1): | |
| r_mean = Y.data.norm(dim=-1).mean().item() | |
| print( | |
| f"[CO-SNE] iter {it:4d}/{self.n_iter} " | |
| f"KL={kl.item():.4f} H={h_loss.item():.6f} " | |
| f"loss={loss.item():.4f} mean_r={r_mean:.4f}" | |
| ) | |
| return project_to_ball(Y.detach()).cpu().numpy() | |
| # --------------------------------------------------------------------------- | |
| # Feature extraction (hyperbolic only; aligned with Evaluate / plot_tsne) | |
| # --------------------------------------------------------------------------- | |
| def preprocess_image_like_forward(image: torch.Tensor) -> torch.Tensor: | |
| if image.ndim == 3: | |
| image = image.unsqueeze(0) | |
| if image.ndim == 4: | |
| if image.shape[1] != 1: | |
| image = image.mean(dim=1, keepdim=True) | |
| elif image.ndim == 5: | |
| if image.shape[1] != 1: | |
| image = image.mean(dim=1, keepdim=True) | |
| else: | |
| raise ValueError(f"Unexpected image ndim={image.ndim}") | |
| return image | |
| def extract_hyperbolic_features(model, dataloader, device, modal_embedding=False): | |
| """Only hyperbolic features (project_img / project_txt).""" | |
| all_hyp_img, all_hyp_txt = [], [] | |
| model.eval() | |
| for batch in dataloader: | |
| imgs = preprocess_image_like_forward( | |
| batch["imgs"].to(device, non_blocking=True) | |
| ) | |
| enc = model.visual_transformer( | |
| imgs, | |
| return_encoded_tokens=True, | |
| modal_embedding=modal_embedding, | |
| modal_indexs=None, | |
| is_condition=False, | |
| ) | |
| if enc.dim() == 2: | |
| image_global = l2norm(enc) | |
| elif enc.dim() == 3: | |
| image_global = l2norm(enc[:, 0, :]) | |
| else: | |
| raise ValueError(f"Unexpected visual shape: {enc.shape}") | |
| text_out = model.text_transformer( | |
| input_ids=batch["caption_ids"].to(device, non_blocking=True), | |
| token_type_ids=batch["token_type_ids"].to(device, non_blocking=True), | |
| attention_mask=batch["attention_mask"].to(device, non_blocking=True), | |
| ) | |
| text_global = text_out[0][:, 0, :] | |
| all_hyp_img.append(model.project_img(image_global).cpu()) | |
| all_hyp_txt.append(model.project_txt(text_global).cpu()) | |
| return torch.cat(all_hyp_img, dim=0), torch.cat(all_hyp_txt, dim=0) | |
| def build_model(args, device: str) -> RADIR: | |
| tokenizer = AutoTokenizer.from_pretrained( | |
| args.model_path, trust_remote_code=True, local_files_only=True | |
| ) | |
| text_model = AutoModel.from_pretrained( | |
| args.model_path, trust_remote_code=True, local_files_only=True | |
| ) | |
| image_encoder = CTViT( | |
| dim=768, | |
| codebook_size=8192, | |
| image_size=args.imsize, | |
| patch_size=16, | |
| temporal_patch_size=10, | |
| spatial_depth=8, | |
| temporal_depth=6, | |
| cls_depth=4, | |
| dim_head=32, | |
| heads=8, | |
| channels=1, | |
| ) | |
| model = RADIR( | |
| tokenizer=tokenizer, | |
| image_encoder=image_encoder, | |
| text_encoder=text_model, | |
| dim_text=768, | |
| dim_image=768, | |
| dim_latent=768, | |
| use_mlm=False, | |
| use_all_token_embeds=False, | |
| ) | |
| print(f"Loading checkpoint: {args.radir_ckpt}") | |
| model.load(args.radir_ckpt) | |
| model.to(device) | |
| return model | |
| def subsample_pair( | |
| hyp_img: torch.Tensor, | |
| hyp_txt: torch.Tensor, | |
| max_samples: int | None, | |
| seed: int, | |
| ): | |
| n = hyp_img.shape[0] | |
| if max_samples is None or max_samples <= 0 or max_samples >= n: | |
| return hyp_img, hyp_txt, np.arange(n) | |
| rng = np.random.default_rng(seed) | |
| idx = np.sort(rng.choice(n, size=max_samples, replace=False)) | |
| return hyp_img[idx], hyp_txt[idx], idx | |
| # --------------------------------------------------------------------------- | |
| # Plotting | |
| # --------------------------------------------------------------------------- | |
| def plot_poincare_disk( | |
| emb_img: np.ndarray, | |
| emb_txt: np.ndarray, | |
| out_path: str, | |
| title: str, | |
| dpi: int, | |
| max_pair_lines: int = 40, | |
| ) -> None: | |
| fig, ax = plt.subplots(figsize=(8, 8)) | |
| circle = plt.Circle((0, 0), 1.0, color="black", fill=False, linewidth=1.2) | |
| ax.add_patch(circle) | |
| # light radial rings to emphasize hierarchy | |
| for r in (0.25, 0.5, 0.75): | |
| ax.add_patch( | |
| plt.Circle((0, 0), r, color="gray", fill=False, linewidth=0.4, alpha=0.4) | |
| ) | |
| n = emb_img.shape[0] | |
| n_lines = min(max_pair_lines, n) | |
| for i in range(n_lines): | |
| ax.plot( | |
| [emb_img[i, 0], emb_txt[i, 0]], | |
| [emb_img[i, 1], emb_txt[i, 1]], | |
| color="gray", | |
| alpha=0.15, | |
| linewidth=0.5, | |
| zorder=0, | |
| ) | |
| ax.scatter( | |
| emb_img[:, 0], | |
| emb_img[:, 1], | |
| c="#2962FF", | |
| marker="o", | |
| s=28, | |
| alpha=0.85, | |
| edgecolors="black", | |
| linewidths=0.25, | |
| label="Image", | |
| zorder=2, | |
| ) | |
| ax.scatter( | |
| emb_txt[:, 0], | |
| emb_txt[:, 1], | |
| c="#D50000", | |
| marker="^", | |
| s=36, | |
| alpha=0.85, | |
| edgecolors="black", | |
| linewidths=0.25, | |
| label="Text", | |
| zorder=3, | |
| ) | |
| # report mean radius (hierarchy signal) | |
| r_img = np.linalg.norm(emb_img, axis=1).mean() | |
| r_txt = np.linalg.norm(emb_txt, axis=1).mean() | |
| ax.set_title( | |
| f"{title}\n" | |
| f"mean Poincaré radius Image={r_img:.3f} Text={r_txt:.3f}" | |
| ) | |
| ax.set_xlim(-1.05, 1.05) | |
| ax.set_ylim(-1.05, 1.05) | |
| ax.set_aspect("equal") | |
| ax.legend(loc="upper right", frameon=True) | |
| ax.axis("off") | |
| fig.tight_layout() | |
| os.makedirs(os.path.dirname(os.path.abspath(out_path)) or ".", exist_ok=True) | |
| fig.savefig(out_path, dpi=dpi, bbox_inches="tight") | |
| # also save pdf sibling if png | |
| if out_path.lower().endswith(".png"): | |
| fig.savefig(out_path[:-4] + ".pdf", dpi=dpi, bbox_inches="tight") | |
| plt.close(fig) | |
| print(f"Saved CO-SNE plot to {out_path}") | |
| print(f" mean radius Image={r_img:.4f}, Text={r_txt:.4f} " | |
| f"(Text closer to origin ⇒ hierarchy)") | |
| def print_radius_stats(name: str, poincare: np.ndarray) -> None: | |
| r = np.linalg.norm(poincare, axis=1) | |
| d0 = np.arctanh(np.clip(r, 0, 1 - 1e-6)) | |
| print( | |
| f" {name:6s} radius mean={r.mean():.4f} std={r.std():.4f} " | |
| f"| d_H(·,0) mean={d0.mean():.4f} std={d0.std():.4f}" | |
| ) | |
| # --------------------------------------------------------------------------- | |
| # CLI | |
| # --------------------------------------------------------------------------- | |
| def parse_args(): | |
| p = argparse.ArgumentParser( | |
| description="CO-SNE visualization of RadIR hyperbolic embeddings" | |
| ) | |
| p.add_argument("--device", default="cuda") | |
| p.add_argument( | |
| "--model_path", | |
| default=os.path.join( | |
| ROOT, "hf_models/microsoft/BiomedVLP-CXR-BERT-specialized" | |
| ), | |
| ) | |
| p.add_argument( | |
| "--radir_ckpt", | |
| default=os.path.join(ROOT, "outputs/radir_final.pt"), | |
| ) | |
| p.add_argument("--dataset_root", default=os.path.join(ROOT, "dataset")) | |
| p.add_argument("--caption_file", default="valid.jsonl") | |
| p.add_argument( | |
| "--cig_matched_dir", | |
| default=os.path.join(ROOT, "dataset/processed/matched"), | |
| ) | |
| p.add_argument("--split", default="valid") | |
| p.add_argument("--batch_size", type=int, default=16) | |
| p.add_argument("--num_workers", type=int, default=8) | |
| p.add_argument("--imsize", type=int, default=224) | |
| p.add_argument("--max_words", type=int, default=128) | |
| p.add_argument("--modal_embedding", action="store_true") | |
| p.add_argument( | |
| "--load_feats", | |
| default=None, | |
| help="Load outputs/tsne_feats.npz (uses hyp_img / hyp_txt / curv only)", | |
| ) | |
| p.add_argument("--save_feats", default=None) | |
| p.add_argument( | |
| "--max_samples", | |
| type=int, | |
| default=3120, | |
| help="Subsample pairs for CO-SNE speed; 0 = all (slow, O(N^2))", | |
| ) | |
| p.add_argument("--perplexity", type=float, default=30.0) | |
| p.add_argument("--n_iter", type=int, default=1000) | |
| p.add_argument("--learning_rate", type=float, default=0.05) | |
| p.add_argument( | |
| "--learning_rate_h", | |
| type=float, | |
| default=5.0, | |
| help="Weight for origin-radius (hierarchy) loss; keep high so radii are preserved", | |
| ) | |
| p.add_argument("--early_exaggeration", type=float, default=4.0) | |
| p.add_argument("--student_t_gamma", type=float, default=0.1) | |
| p.add_argument("--seed", type=int, default=42) | |
| p.add_argument("--max_pair_lines", type=int, default=40) | |
| p.add_argument( | |
| "--out", | |
| default=os.path.join(ROOT, "outputs/cosne_radir.png"), | |
| ) | |
| p.add_argument( | |
| "--save_emb", | |
| default=None, | |
| help="Optional .npz to cache 2D CO-SNE Poincaré coords", | |
| ) | |
| p.add_argument("--dpi", type=int, default=150) | |
| p.add_argument( | |
| "--cosne_device", | |
| default="cpu", | |
| help="Device for CO-SNE optimization (cpu is usually fine for N~1k)", | |
| ) | |
| return p.parse_args() | |
| def main(): | |
| args = parse_args() | |
| device = args.device | |
| if args.load_feats: | |
| print(f"Loading cached features: {args.load_feats}") | |
| data = np.load(args.load_feats, allow_pickle=True) | |
| if "hyp_img" not in data.files or "hyp_txt" not in data.files: | |
| raise KeyError( | |
| f"{args.load_feats} must contain hyp_img / hyp_txt " | |
| "(hyperbolic features). Re-extract with plot_tsne.py --save_feats." | |
| ) | |
| hyp_img = torch.from_numpy(data["hyp_img"]) | |
| hyp_txt = torch.from_numpy(data["hyp_txt"]) | |
| curv = float(data["curv"]) if "curv" in data.files else 1.0 | |
| print(f" hyp_img={tuple(hyp_img.shape)}, curv={curv:.6f}") | |
| else: | |
| model = build_model(args, device) | |
| curv = float(model.get_curv().detach().cpu().item()) | |
| print(f"Curvature = {curv:.6f}") | |
| dataset = HyCoClipMimicBoxDataset_JPG( | |
| split=args.split, | |
| tokenizer=model.tokenizer, | |
| dataset_root=args.dataset_root, | |
| cig_matched_dir=args.cig_matched_dir, | |
| caption_file=args.caption_file, | |
| imsize=args.imsize, | |
| max_words=args.max_words, | |
| transform=None, | |
| ) | |
| dataloader = DataLoader( | |
| dataset, | |
| batch_size=args.batch_size, | |
| shuffle=False, | |
| num_workers=args.num_workers, | |
| pin_memory=True, | |
| collate_fn=collate_hmclip_box, | |
| ) | |
| print(f"Extracting hyperbolic features from {len(dataset)} samples ...") | |
| hyp_img, hyp_txt = extract_hyperbolic_features( | |
| model, dataloader, device, modal_embedding=args.modal_embedding | |
| ) | |
| if args.save_feats: | |
| os.makedirs( | |
| os.path.dirname(os.path.abspath(args.save_feats)) or ".", | |
| exist_ok=True, | |
| ) | |
| np.savez( | |
| args.save_feats, | |
| hyp_img=hyp_img.numpy(), | |
| hyp_txt=hyp_txt.numpy(), | |
| curv=curv, | |
| ) | |
| print(f"Saved hyperbolic features to {args.save_feats}") | |
| max_samples = None if args.max_samples <= 0 else args.max_samples | |
| hyp_img, hyp_txt, idx = subsample_pair( | |
| hyp_img, hyp_txt, max_samples, args.seed | |
| ) | |
| print(f"Using {hyp_img.shape[0]} pairs for CO-SNE") | |
| # Lorentz space → Poincaré (unit disk) | |
| with torch.no_grad(): | |
| poi_img = lorentz_space_to_poincare(hyp_img.float(), curv).numpy() | |
| poi_txt = lorentz_space_to_poincare(hyp_txt.float(), curv).numpy() | |
| print("High-dim Poincaré radius / origin-distance stats:") | |
| print_radius_stats("Image", poi_img) | |
| print_radius_stats("Text", poi_txt) | |
| # Joint CO-SNE on Image + Text (preserves cross-modal + hierarchy) | |
| X = np.concatenate([poi_img, poi_txt], axis=0) | |
| cosne = COSNE( | |
| n_components=2, | |
| perplexity=args.perplexity, | |
| n_iter=args.n_iter, | |
| learning_rate=args.learning_rate, | |
| learning_rate_h=args.learning_rate_h, | |
| early_exaggeration=args.early_exaggeration, | |
| student_t_gamma=args.student_t_gamma, | |
| seed=args.seed, | |
| device=args.cosne_device, | |
| verbose=True, | |
| ) | |
| Y = cosne.fit_transform(X) | |
| n = poi_img.shape[0] | |
| emb_img, emb_txt = Y[:n], Y[n:] | |
| print("Low-dim (CO-SNE) Poincaré radius / origin-distance stats:") | |
| print_radius_stats("Image", emb_img) | |
| print_radius_stats("Text", emb_txt) | |
| if args.save_emb: | |
| os.makedirs( | |
| os.path.dirname(os.path.abspath(args.save_emb)) or ".", exist_ok=True | |
| ) | |
| np.savez( | |
| args.save_emb, | |
| emb_img=emb_img, | |
| emb_txt=emb_txt, | |
| sample_idx=idx, | |
| curv=curv, | |
| ) | |
| print(f"Saved 2D embeddings to {args.save_emb}") | |
| ckpt = os.path.basename(args.radir_ckpt) if not args.load_feats else "cached" | |
| title = f"CO-SNE (hyperbolic) | RadIR | {ckpt} | n={n}" | |
| plot_poincare_disk( | |
| emb_img, | |
| emb_txt, | |
| args.out, | |
| title=title, | |
| dpi=args.dpi, | |
| max_pair_lines=args.max_pair_lines, | |
| ) | |
| if __name__ == "__main__": | |
| main() | |