HMCLIP / scripts /plot_cosne.py
SteveZeyuZhang's picture
Add source files and streaming backup restore instructions
b91f38a
Raw History Blame Contribute Delete
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
@torch.no_grad()
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()