|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import math, time, json, os |
|
|
from pathlib import Path |
|
|
from tqdm.auto import tqdm |
|
|
import torch |
|
|
import torch.nn as nn |
|
|
import torch.nn.functional as F |
|
|
|
|
|
HF_REPO = "AbstractPhil/grid-geometric-classifier-proto" |
|
|
CKPT_DIR = Path("./checkpoints") |
|
|
CC_CKPT_DIR = Path("./cc_checkpoints") |
|
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
torch.backends.cuda.matmul.allow_tf32 = True |
|
|
torch.backends.cudnn.allow_tf32 = True |
|
|
torch.backends.cudnn.benchmark = True |
|
|
|
|
|
use_amp = device.type == "cuda" |
|
|
amp_dtype = torch.bfloat16 if (device.type == "cuda" and |
|
|
torch.cuda.is_bf16_supported()) else torch.float16 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
SHAPE_DESCRIPTIONS = { |
|
|
"point": "A zero-dimensional geometric primitive occupying a single discrete location in three-dimensional space with no extent along any axis.", |
|
|
"line_x": "A one-dimensional line segment extending along the horizontal x-axis, connecting two endpoints with uniform spacing between occupied voxels.", |
|
|
"line_y": "A one-dimensional line segment extending along the vertical y-axis, a straight structure rising upward through the grid.", |
|
|
"line_z": "A one-dimensional line segment extending along the depth z-axis, projecting straight backward through the voxel grid.", |
|
|
"line_diag": "A one-dimensional diagonal line segment cutting across multiple axes simultaneously, connecting opposite corners of the grid.", |
|
|
"cross": "Two perpendicular line segments intersecting at their midpoints forming a plus-shaped cross pattern in a single plane.", |
|
|
"l_shape": "Two connected line segments meeting at a right angle to form an L-shaped corner, like two edges of a rectangle.", |
|
|
"collinear": "Three or more points arranged along a single straight line with equal spacing, demonstrating perfect linear alignment.", |
|
|
"triangle_xy": "A flat triangular outline formed by three connected edges lying in the horizontal xy-plane, the simplest polygon.", |
|
|
"triangle_xz": "A flat triangular outline formed by three connected edges lying in the vertical xz-plane, a triangle standing upright.", |
|
|
"triangle_3d": "A triangular outline with vertices at different heights, forming a non-planar triangle tilted in three-dimensional space.", |
|
|
"square_xy": "A square outline formed by four equal edges in the xy-plane, a regular quadrilateral with right angles at each corner.", |
|
|
"square_xz": "A square outline formed by four equal edges in the xz-plane, a square standing vertically like a window frame.", |
|
|
"rectangle": "A rectangular outline with two pairs of parallel edges of different lengths, wider than it is tall.", |
|
|
"coplanar": "A set of points all lying in the same plane but not forming a regular polygon, a scattered planar arrangement.", |
|
|
"plane": "A solid flat surface filling an entire plane with occupied voxels, a two-dimensional sheet extending across the grid.", |
|
|
"tetrahedron": "A three-dimensional simplex with four triangular faces meeting at four vertices and six edges, the simplest polyhedron.", |
|
|
"pyramid": "A solid with a square base and four triangular faces converging to a single apex point above the base center.", |
|
|
"pentachoron": "A four-dimensional simplex projected into three dimensions, consisting of five tetrahedral cells sharing faces.", |
|
|
"cube": "A regular hexahedron with six identical square faces, twelve edges, and eight vertices forming a perfect box shape.", |
|
|
"cuboid": "A rectangular box with six rectangular faces, similar to a cube but with at least one pair of edges longer than the others.", |
|
|
"triangular_prism": "A solid with two parallel triangular faces connected by three rectangular faces, like a tent or Toblerone shape.", |
|
|
"octahedron": "A regular polyhedron with eight equilateral triangular faces, twelve edges, and six vertices, like two pyramids base-to-base.", |
|
|
"arc": "A curved one-dimensional segment forming part of a circle, a smooth bend connecting two endpoints along a circular path.", |
|
|
"helix": "A three-dimensional spiral curve that winds around a central axis while advancing along it, like a corkscrew or spring.", |
|
|
"circle": "A closed curved outline where every point is equidistant from the center, forming a perfect round ring in a plane.", |
|
|
"ellipse": "A closed curved outline forming an elongated circle, an oval shape with two focal points and varying curvature.", |
|
|
"disc": "A solid filled circular region, a flat round plate occupying all voxels within a circular boundary in a plane.", |
|
|
"sphere": "A perfectly round three-dimensional solid where every surface point is equidistant from the center, fully filled inside.", |
|
|
"hemisphere": "Half of a sphere cut along a great circle, a dome shape with a flat base and a convex curved upper surface.", |
|
|
"cylinder": "A solid with two parallel circular faces connected by a curved rectangular surface, like a can or pillar.", |
|
|
"cone": "A solid tapering smoothly from a circular base to a single apex point, with a curved surface of decreasing radius.", |
|
|
"capsule": "A cylinder capped with hemispheres at both ends, a smooth elongated pill shape with no sharp edges.", |
|
|
"torus": "A donut-shaped solid formed by revolving a circle around an external axis, with a hole through the center.", |
|
|
"shell": "A hollow spherical surface with no interior fill, an empty ball where only the outer boundary layer is occupied.", |
|
|
"tube": "A hollow cylindrical surface with no interior fill, an empty pipe where only the curved wall is occupied.", |
|
|
"bowl": "A concave open surface curving inward like a dish, the bottom half of a hollow sphere with the opening facing up.", |
|
|
"saddle": "A hyperbolic surface that curves upward along one axis and downward along the perpendicular axis, like a horse saddle.", |
|
|
} |
|
|
assert set(SHAPE_DESCRIPTIONS.keys()) == set(CLASS_NAMES), "Description/class mismatch!" |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class QwenEmbeddingExtractor: |
|
|
MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct" |
|
|
HIDDEN_DIM = 1536 |
|
|
|
|
|
def __init__(self, device="cuda"): |
|
|
self.device = device |
|
|
self.model = None |
|
|
self.tokenizer = None |
|
|
|
|
|
def load_model(self): |
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
print(f"Loading {self.MODEL_NAME}...") |
|
|
self.tokenizer = AutoTokenizer.from_pretrained(self.MODEL_NAME, trust_remote_code=True) |
|
|
self.model = AutoModelForCausalLM.from_pretrained( |
|
|
self.MODEL_NAME, dtype=torch.float16, |
|
|
device_map=self.device, trust_remote_code=True) |
|
|
self.model.eval() |
|
|
print(f"Qwen loaded: {self.HIDDEN_DIM}-dim hidden states") |
|
|
|
|
|
def _build_encode_prompt(self, description): |
|
|
messages = [ |
|
|
{"role": "system", "content": "You are a geometric shape analyst."}, |
|
|
{"role": "user", "content": f"Analyze this shape: {description}"}, |
|
|
] |
|
|
return self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) |
|
|
|
|
|
@torch.no_grad() |
|
|
def extract_embedding(self, text): |
|
|
prompt = self._build_encode_prompt(text) |
|
|
inputs = self.tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True).to(self.device) |
|
|
outputs = self.model(**inputs, output_hidden_states=True) |
|
|
hidden = outputs.hidden_states[-1] |
|
|
return hidden.mean(dim=1).squeeze(0).float() |
|
|
|
|
|
def cache_all_embeddings(self, class_names): |
|
|
print(f"Extracting embeddings for {len(class_names)} classes...") |
|
|
embeddings = {} |
|
|
for name in class_names: |
|
|
embeddings[name] = self.extract_embedding(SHAPE_DESCRIPTIONS[name]) |
|
|
emb_tensor = torch.stack([embeddings[n] for n in class_names]) |
|
|
normed = F.normalize(emb_tensor, dim=-1) |
|
|
sim = normed @ normed.T |
|
|
mean_sim = (sim.sum() - sim.trace()) / (len(class_names) * (len(class_names) - 1)) |
|
|
print(f"Cached: {emb_tensor.shape} | mean cross-class sim: {mean_sim:.4f}") |
|
|
return emb_tensor |
|
|
|
|
|
def unload(self): |
|
|
del self.model, self.tokenizer |
|
|
self.model = self.tokenizer = None |
|
|
torch.cuda.empty_cache() |
|
|
print("Qwen unloaded") |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class TextProjection(nn.Module): |
|
|
def __init__(self, text_dim=1536, latent_dim=256): |
|
|
super().__init__() |
|
|
self.proj = nn.Sequential( |
|
|
nn.Linear(text_dim, latent_dim * 2), nn.GELU(), |
|
|
nn.Linear(latent_dim * 2, latent_dim), nn.GELU(), |
|
|
nn.Linear(latent_dim, latent_dim)) |
|
|
self.norm = nn.LayerNorm(latent_dim) |
|
|
def forward(self, x): return self.norm(self.proj(x)) |
|
|
|
|
|
class VoxelProjection(nn.Module): |
|
|
def __init__(self, voxel_dim=645, latent_dim=256): |
|
|
super().__init__() |
|
|
self.proj = nn.Sequential( |
|
|
nn.Linear(voxel_dim, latent_dim * 2), nn.GELU(), |
|
|
nn.Linear(latent_dim * 2, latent_dim), nn.GELU(), |
|
|
nn.Linear(latent_dim, latent_dim)) |
|
|
self.norm = nn.LayerNorm(latent_dim) |
|
|
def forward(self, x): return self.norm(self.proj(x)) |
|
|
|
|
|
|
|
|
class CrossContrastModel(nn.Module): |
|
|
def __init__(self, text_dim=1536, voxel_dim=645, latent_dim=256, |
|
|
n_classes=38, temperature=0.07): |
|
|
super().__init__() |
|
|
self.text_proj = TextProjection(text_dim, latent_dim) |
|
|
self.voxel_proj = VoxelProjection(voxel_dim, latent_dim) |
|
|
self.log_temperature = nn.Parameter(torch.tensor(math.log(1.0 / temperature))) |
|
|
|
|
|
@property |
|
|
def temperature(self): |
|
|
return torch.exp(-self.log_temperature) |
|
|
|
|
|
def forward(self, voxel_features, class_labels, text_embeddings_table): |
|
|
text_emb = text_embeddings_table[class_labels] |
|
|
z_text = F.normalize(self.text_proj(text_emb), dim=-1) |
|
|
z_voxel = F.normalize(self.voxel_proj(voxel_features), dim=-1) |
|
|
|
|
|
temp = self.temperature |
|
|
logits_v2t = z_voxel @ z_text.T / temp |
|
|
logits_t2v = z_text @ z_voxel.T / temp |
|
|
|
|
|
labels_matrix = (class_labels.unsqueeze(0) == class_labels.unsqueeze(1)).float() |
|
|
labels_matrix = labels_matrix / labels_matrix.sum(dim=1, keepdim=True).clamp(min=1) |
|
|
|
|
|
loss_v2t = (-labels_matrix * F.log_softmax(logits_v2t, dim=1)).sum(dim=1).mean() |
|
|
loss_t2v = (-labels_matrix * F.log_softmax(logits_t2v, dim=1)).sum(dim=1).mean() |
|
|
loss = (loss_v2t + loss_t2v) / 2.0 |
|
|
|
|
|
with torch.no_grad(): |
|
|
v2t_preds = logits_v2t.argmax(dim=1) |
|
|
pred_classes = class_labels[v2t_preds] |
|
|
acc = (pred_classes == class_labels).float().mean() |
|
|
pos_sim = (z_voxel * z_text).sum(dim=-1).mean() |
|
|
neg_mask = ~(class_labels.unsqueeze(0) == class_labels.unsqueeze(1)) |
|
|
neg_sim = (z_voxel @ z_text.T)[neg_mask].mean() if neg_mask.any() else torch.tensor(0.0) |
|
|
|
|
|
return loss, {"acc": acc.item(), "pos_sim": pos_sim.item(), |
|
|
"neg_sim": neg_sim.item(), "temperature": temp.item()} |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def get_hf_token(): |
|
|
try: |
|
|
from google.colab import userdata |
|
|
return userdata.get('HF_TOKEN') |
|
|
except Exception: |
|
|
return os.environ.get('HF_TOKEN') |
|
|
|
|
|
|
|
|
def load_geo_model(ckpt_dir=CKPT_DIR): |
|
|
"""Load trained GeometricShapeClassifier from Cell 3 checkpoint.""" |
|
|
latest = ckpt_dir / "latest.pt" |
|
|
if not latest.exists(): |
|
|
raise FileNotFoundError( |
|
|
f"No checkpoint at {latest}. Run Cell 3 first to train the classifier.") |
|
|
print(f"Loading geo classifier from {latest}...") |
|
|
ckpt = torch.load(latest, weights_only=False, map_location=device) |
|
|
geo = GeometricShapeClassifier().to(device) |
|
|
geo.load_state_dict(ckpt["model_state_dict"]) |
|
|
geo.eval() |
|
|
for p in geo.parameters(): |
|
|
p.requires_grad = False |
|
|
print(f"Loaded: epoch {ckpt['epoch']}, val_acc={ckpt['best_val_acc']:.4f}, " |
|
|
f"{sum(p.numel() for p in geo.parameters()):,} params (frozen)") |
|
|
return geo |
|
|
|
|
|
|
|
|
@torch.no_grad() |
|
|
def extract_voxel_features(geo_model, grid): |
|
|
"""Extract pre-classifier features using model.forward()['features'].""" |
|
|
with torch.amp.autocast('cuda', enabled=use_amp, dtype=amp_dtype): |
|
|
out = geo_model(grid) |
|
|
return out["features"].float() |
|
|
|
|
|
|
|
|
def save_cc_checkpoint(cc_model, cc_opt, cc_sched, epoch, best_acc, ckpt_dir=CC_CKPT_DIR): |
|
|
ckpt_dir.mkdir(parents=True, exist_ok=True) |
|
|
ckpt = { |
|
|
"epoch": epoch, |
|
|
"best_cc_acc": best_acc, |
|
|
"cc_model_state_dict": cc_model.state_dict(), |
|
|
"cc_optimizer_state_dict": cc_opt.state_dict(), |
|
|
"cc_scheduler_state_dict": cc_sched.state_dict(), |
|
|
} |
|
|
torch.save(ckpt, ckpt_dir / "latest.pt") |
|
|
|
|
|
|
|
|
def load_cc_checkpoint(cc_model, cc_opt, cc_sched, ckpt_dir=CC_CKPT_DIR): |
|
|
latest = ckpt_dir / "latest.pt" |
|
|
if not latest.exists(): |
|
|
return 0, 0.0 |
|
|
print(f"Resuming CC from {latest}...") |
|
|
ckpt = torch.load(latest, weights_only=False) |
|
|
cc_model.load_state_dict(ckpt["cc_model_state_dict"]) |
|
|
cc_opt.load_state_dict(ckpt["cc_optimizer_state_dict"]) |
|
|
cc_sched.load_state_dict(ckpt["cc_scheduler_state_dict"]) |
|
|
start = ckpt["epoch"] + 1 |
|
|
best = ckpt["best_cc_acc"] |
|
|
print(f"Resumed: epoch {start}, best_cc_acc={best:.4f}") |
|
|
return start, best |
|
|
|
|
|
|
|
|
def upload_cc_to_hf(cc_model, best_acc, epoch, token, reason="periodic", |
|
|
text_dim_=None, voxel_dim_=None, latent_dim_=None): |
|
|
"""Upload crosscontrast weights + config to HF. Called mid-training and at end.""" |
|
|
if not token: |
|
|
return |
|
|
try: |
|
|
from huggingface_hub import HfApi, create_repo |
|
|
from safetensors.torch import save_file as st_save |
|
|
|
|
|
staging = Path("./hf_staging/crosscontrast") |
|
|
staging.mkdir(parents=True, exist_ok=True) |
|
|
|
|
|
st_save(cc_model.text_proj.state_dict(), str(staging / "text_proj.safetensors")) |
|
|
st_save(cc_model.voxel_proj.state_dict(), str(staging / "voxel_proj.safetensors")) |
|
|
st_save({"log_temperature": cc_model.log_temperature.data.unsqueeze(0)}, |
|
|
str(staging / "temperature.safetensors")) |
|
|
|
|
|
|
|
|
if text_dim_ and voxel_dim_ and latent_dim_: |
|
|
cfg = { |
|
|
"model_type": "CrossContrastModel", |
|
|
"text_dim": text_dim_, "voxel_dim": voxel_dim_, |
|
|
"latent_dim": latent_dim_, |
|
|
"best_val_accuracy": best_acc, |
|
|
"epoch": epoch + 1, "upload_reason": reason, |
|
|
"temperature": cc_model.temperature.item(), |
|
|
} |
|
|
with open(staging / "config.json", "w") as f: |
|
|
json.dump(cfg, f, indent=2) |
|
|
|
|
|
api = HfApi(token=token) |
|
|
create_repo(HF_REPO, token=token, exist_ok=True) |
|
|
api.upload_folder( |
|
|
folder_path=str(staging), repo_id=HF_REPO, |
|
|
path_in_repo="crosscontrast", token=token, |
|
|
commit_message=f"crosscontrast ep{epoch+1} | acc={best_acc:.4f} | {reason}") |
|
|
tqdm.write(f" ✓ HF upload ({reason}): ep{epoch+1} acc={best_acc:.4f}") |
|
|
except Exception as e: |
|
|
tqdm.write(f" ⚠ HF upload failed: {e}") |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
print("=" * 70) |
|
|
print("Phase 1: Qwen Embeddings") |
|
|
print("=" * 70) |
|
|
|
|
|
cache_path = Path("qwen_geo_cache.pt") |
|
|
if cache_path.exists(): |
|
|
_cache = torch.load(cache_path, map_location="cpu", weights_only=True) |
|
|
text_embeddings = _cache["embeddings"] |
|
|
print(f"Loaded cached: {text_embeddings.shape}") |
|
|
else: |
|
|
extractor = QwenEmbeddingExtractor(device=str(device)) |
|
|
extractor.load_model() |
|
|
text_embeddings = extractor.cache_all_embeddings(CLASS_NAMES) |
|
|
torch.save({"embeddings": text_embeddings.cpu(), "class_names": CLASS_NAMES}, cache_path) |
|
|
extractor.unload() |
|
|
|
|
|
text_embeddings = text_embeddings.to(device) |
|
|
text_dim = text_embeddings.shape[1] |
|
|
|
|
|
|
|
|
hf_token = get_hf_token() |
|
|
qwen_staging = Path("./hf_staging/qwen_embeddings") |
|
|
qwen_staging.mkdir(parents=True, exist_ok=True) |
|
|
|
|
|
qwen_config = { |
|
|
"model_name": QwenEmbeddingExtractor.MODEL_NAME, |
|
|
"hidden_dim": QwenEmbeddingExtractor.HIDDEN_DIM, |
|
|
"extraction_method": "mean_pool_last_layer", |
|
|
"prompt_style": "2shot_geometric", |
|
|
"num_classes": NUM_CLASSES, |
|
|
"class_names": CLASS_NAMES, |
|
|
"embedding_shape": list(text_embeddings.shape), |
|
|
} |
|
|
with open(qwen_staging / "config.json", "w") as f: |
|
|
json.dump(qwen_config, f, indent=2) |
|
|
with open(qwen_staging / "descriptions.json", "w") as f: |
|
|
json.dump(SHAPE_DESCRIPTIONS, f, indent=2) |
|
|
|
|
|
try: |
|
|
from safetensors.torch import save_file as st_save |
|
|
st_save({"embeddings": text_embeddings.cpu()}, str(qwen_staging / "embeddings.safetensors")) |
|
|
except ImportError: |
|
|
torch.save({"embeddings": text_embeddings.cpu()}, qwen_staging / "embeddings.pt") |
|
|
|
|
|
if hf_token: |
|
|
try: |
|
|
from huggingface_hub import HfApi, create_repo |
|
|
api = HfApi(token=hf_token) |
|
|
create_repo(HF_REPO, token=hf_token, exist_ok=True) |
|
|
api.upload_folder( |
|
|
folder_path=str(qwen_staging), repo_id=HF_REPO, |
|
|
path_in_repo="qwen_embeddings", token=hf_token, |
|
|
commit_message=f"qwen_embeddings | {QwenEmbeddingExtractor.MODEL_NAME} | {NUM_CLASSES} classes") |
|
|
print(f"Uploaded: https://huggingface.co/{HF_REPO}/tree/main/qwen_embeddings") |
|
|
except Exception as e: |
|
|
print(f"Upload failed: {e}") |
|
|
else: |
|
|
print("No HF_TOKEN — saved locally") |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
print("\n" + "=" * 70) |
|
|
print("Phase 2: Geo Model + Voxel Data") |
|
|
print("=" * 70) |
|
|
|
|
|
geo_model = load_geo_model() |
|
|
|
|
|
CC_SAMPLES = 500000 |
|
|
CC_BATCH = 4096 |
|
|
CC_EPOCHS = 40 |
|
|
CC_LR = 2e-3 |
|
|
CC_LATENT = 256 |
|
|
|
|
|
|
|
|
DATASET_PATH = Path("./cached_dataset.pt") |
|
|
if not DATASET_PATH.exists(): |
|
|
raise FileNotFoundError("No cached_dataset.pt — run Cell 3 first.") |
|
|
|
|
|
print(f"Loading dataset from {DATASET_PATH}...") |
|
|
_cached = torch.load(DATASET_PATH, weights_only=True) |
|
|
cc_train_ds = ShapeDataset.__new__(ShapeDataset) |
|
|
cc_val_ds = ShapeDataset.__new__(ShapeDataset) |
|
|
for k in ["grids", "labels", "dim_conf", "peak_dim", "volume", "cm_det", "is_curved", "curvature"]: |
|
|
setattr(cc_train_ds, k, _cached["train"][k]) |
|
|
setattr(cc_val_ds, k, _cached["val"][k]) |
|
|
print(f"Loaded {len(cc_train_ds)} train + {len(cc_val_ds)} val (from Cell 3 cache)") |
|
|
|
|
|
cc_train_loader = torch.utils.data.DataLoader( |
|
|
cc_train_ds, batch_size=CC_BATCH, shuffle=True, |
|
|
num_workers=4, pin_memory=True, persistent_workers=True) |
|
|
cc_val_loader = torch.utils.data.DataLoader( |
|
|
cc_val_ds, batch_size=CC_BATCH, shuffle=False, |
|
|
num_workers=4, pin_memory=True, persistent_workers=True) |
|
|
|
|
|
|
|
|
with torch.no_grad(): |
|
|
_dummy = torch.zeros(1, GS, GS, GS, device=device) |
|
|
voxel_dim = extract_voxel_features(geo_model, _dummy).shape[1] |
|
|
print(f"Voxel dim: {voxel_dim} | Text dim: {text_dim}") |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
print("\n" + "=" * 70) |
|
|
print("Phase 3: Cross-Contrast Training") |
|
|
print("=" * 70) |
|
|
|
|
|
cc_model = CrossContrastModel( |
|
|
text_dim=text_dim, voxel_dim=voxel_dim, |
|
|
latent_dim=CC_LATENT, n_classes=NUM_CLASSES |
|
|
).to(device) |
|
|
cc_params = sum(p.numel() for p in cc_model.parameters()) |
|
|
print(f"CrossContrast: {cc_params:,} params | latent={CC_LATENT}") |
|
|
|
|
|
cc_opt = torch.optim.AdamW(cc_model.parameters(), lr=CC_LR, weight_decay=1e-4) |
|
|
_warmup = 3 |
|
|
def _cc_lr(ep): |
|
|
if ep < _warmup: return (ep + 1) / _warmup |
|
|
return 0.5 * (1 + math.cos(math.pi * (ep - _warmup) / (CC_EPOCHS - _warmup))) |
|
|
cc_sched = torch.optim.lr_scheduler.LambdaLR(cc_opt, _cc_lr) |
|
|
|
|
|
|
|
|
cc_start, best_cc_acc = load_cc_checkpoint(cc_model, cc_opt, cc_sched) |
|
|
|
|
|
t_start = time.time() |
|
|
epoch_bar = tqdm(range(cc_start, CC_EPOCHS), desc="Training", unit="ep") |
|
|
|
|
|
for epoch in epoch_bar: |
|
|
t0 = time.time() |
|
|
cc_model.train() |
|
|
tot_loss, tot_acc, nb = 0, 0, 0 |
|
|
|
|
|
batch_bar = tqdm(cc_train_loader, desc=f"Ep {epoch+1}/{CC_EPOCHS} train", |
|
|
leave=False, unit="batch") |
|
|
for grid, label, *_ in batch_bar: |
|
|
grid = grid.to(device, non_blocking=True) |
|
|
label = label.to(device, non_blocking=True) |
|
|
vf = extract_voxel_features(geo_model, grid) |
|
|
|
|
|
cc_opt.zero_grad(set_to_none=True) |
|
|
with torch.amp.autocast('cuda', enabled=use_amp, dtype=amp_dtype): |
|
|
loss, metrics = cc_model(vf, label, text_embeddings) |
|
|
|
|
|
loss.backward() |
|
|
torch.nn.utils.clip_grad_norm_(cc_model.parameters(), 1.0) |
|
|
cc_opt.step() |
|
|
tot_loss += loss.item(); tot_acc += metrics["acc"]; nb += 1 |
|
|
batch_bar.set_postfix(loss=f"{loss.item():.4f}", acc=f"{metrics['acc']:.3f}") |
|
|
|
|
|
cc_sched.step() |
|
|
|
|
|
cc_model.eval() |
|
|
vl, va, vps, vns, vnb = 0, 0, 0, 0, 0 |
|
|
with torch.no_grad(), torch.amp.autocast('cuda', enabled=use_amp, dtype=amp_dtype): |
|
|
for grid, label, *_ in tqdm(cc_val_loader, desc=f"Ep {epoch+1}/{CC_EPOCHS} val", |
|
|
leave=False, unit="batch"): |
|
|
grid = grid.to(device, non_blocking=True) |
|
|
label = label.to(device, non_blocking=True) |
|
|
vf = extract_voxel_features(geo_model, grid) |
|
|
loss, m = cc_model(vf, label, text_embeddings) |
|
|
vl += loss.item(); va += m["acc"]; vps += m["pos_sim"]; vns += m["neg_sim"]; vnb += 1 |
|
|
|
|
|
tl = tot_loss/max(nb,1); ta = tot_acc/max(nb,1) |
|
|
vl_ = vl/max(vnb,1); va_ = va/max(vnb,1) |
|
|
ps = vps/max(vnb,1); ns = vns/max(vnb,1) |
|
|
temp = cc_model.temperature.item() |
|
|
dt = time.time() - t0 |
|
|
mk = " *" if va_ > best_cc_acc else "" |
|
|
if va_ > best_cc_acc: best_cc_acc = va_ |
|
|
|
|
|
save_cc_checkpoint(cc_model, cc_opt, cc_sched, epoch, best_cc_acc) |
|
|
|
|
|
|
|
|
is_new_best = mk == " *" |
|
|
periodic = (epoch + 1) % 10 == 0 |
|
|
if is_new_best: |
|
|
upload_cc_to_hf(cc_model, best_cc_acc, epoch, hf_token, reason="new_best", |
|
|
text_dim_=text_dim, voxel_dim_=voxel_dim, latent_dim_=CC_LATENT) |
|
|
elif periodic: |
|
|
upload_cc_to_hf(cc_model, best_cc_acc, epoch, hf_token, reason="periodic", |
|
|
text_dim_=text_dim, voxel_dim_=voxel_dim, latent_dim_=CC_LATENT) |
|
|
|
|
|
epoch_bar.set_postfix(loss=f"{vl_:.4f}", acc=f"{va_:.3f}", best=f"{best_cc_acc:.3f}", |
|
|
tau=f"{temp:.4f}") |
|
|
|
|
|
if (epoch+1) % 5 == 0 or epoch == cc_start or mk: |
|
|
tqdm.write(f"Ep {epoch+1:3d}/{CC_EPOCHS} [{dt:.1f}s] | " |
|
|
f"loss {tl:.4f}/{vl_:.4f} | acc {ta:.3f}/{va_:.3f} | " |
|
|
f"pos {ps:.3f} neg {ns:.3f} | τ {temp:.4f}{mk}") |
|
|
|
|
|
if epoch == cc_start and device.type == "cuda": |
|
|
tqdm.write(f"VRAM peak: {torch.cuda.max_memory_allocated()/1e9:.2f}GB") |
|
|
|
|
|
tt = time.time() - t_start |
|
|
print(f"\nDone in {tt:.0f}s ({tt/60:.1f}min) | Best acc: {best_cc_acc:.4f}") |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
print("\n" + "=" * 70) |
|
|
print("Per-Class Alignment") |
|
|
print("=" * 70) |
|
|
|
|
|
cc_model.eval() |
|
|
cls_vox = {n: [] for n in CLASS_NAMES} |
|
|
with torch.no_grad(): |
|
|
text_proj_all = F.normalize(cc_model.text_proj(text_embeddings), dim=-1) |
|
|
|
|
|
with torch.no_grad(), torch.amp.autocast('cuda', enabled=use_amp, dtype=amp_dtype): |
|
|
for grid, label, *_ in cc_val_loader: |
|
|
grid = grid.to(device, non_blocking=True) |
|
|
label = label.to(device, non_blocking=True) |
|
|
vf = extract_voxel_features(geo_model, grid) |
|
|
zv = F.normalize(cc_model.voxel_proj(vf), dim=-1) |
|
|
for k in range(len(label)): |
|
|
cls_vox[CLASS_NAMES[label[k].item()]].append(zv[k].cpu()) |
|
|
|
|
|
print(f"\n{'Class':22s} | {'Align':>6s} | {'N':>5s} | {'Nearest':22s} | {'OK':>3s}") |
|
|
print("-" * 70) |
|
|
correct = 0; total_c = 0 |
|
|
for name in CLASS_NAMES: |
|
|
if not cls_vox[name]: continue |
|
|
mv = F.normalize(torch.stack(cls_vox[name]).mean(dim=0), dim=-1) |
|
|
own = text_proj_all[CLASS_NAMES.index(name)].cpu() |
|
|
align = (mv * own).sum().item() |
|
|
sims = (mv.unsqueeze(0) @ text_proj_all.cpu().T).squeeze(0) |
|
|
ni = sims.argmax().item(); nn_ = CLASS_NAMES[ni] |
|
|
ok = "Y" if nn_ == name else f"X->{nn_}" |
|
|
if nn_ == name: correct += 1 |
|
|
total_c += 1 |
|
|
print(f" {name:20s} | {align:.4f} | {len(cls_vox[name]):5d} | {nn_:22s} | {ok}") |
|
|
print(f"\nNearest-text accuracy: {correct}/{total_c} = {correct/max(total_c,1):.1%}") |
|
|
|
|
|
print("\nTop 10 Text-Space Confusions:") |
|
|
with torch.no_grad(): |
|
|
sim = (text_proj_all @ text_proj_all.T).cpu().numpy() |
|
|
import numpy as np |
|
|
confusions = [] |
|
|
for i in range(len(CLASS_NAMES)): |
|
|
for j in range(i+1, len(CLASS_NAMES)): |
|
|
confusions.append((CLASS_NAMES[i], CLASS_NAMES[j], sim[i,j])) |
|
|
confusions.sort(key=lambda x: x[2], reverse=True) |
|
|
for a, b, s in confusions[:10]: |
|
|
print(f" {a:20s} <-> {b:20s} | {s:.4f}") |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
print("\n" + "=" * 70) |
|
|
print("Saving crosscontrast/ to HuggingFace") |
|
|
print("=" * 70) |
|
|
|
|
|
cc_staging = Path("./hf_staging/crosscontrast") |
|
|
cc_staging.mkdir(parents=True, exist_ok=True) |
|
|
|
|
|
|
|
|
cc_arch = { |
|
|
"model_type": "CrossContrastModel", |
|
|
"text_dim": text_dim, |
|
|
"voxel_dim": voxel_dim, |
|
|
"latent_dim": CC_LATENT, |
|
|
"n_classes": NUM_CLASSES, |
|
|
"text_proj_layers": [text_dim, CC_LATENT * 2, CC_LATENT, CC_LATENT], |
|
|
"voxel_proj_layers": [voxel_dim, CC_LATENT * 2, CC_LATENT, CC_LATENT], |
|
|
"activation": "GELU", |
|
|
"normalization": "LayerNorm", |
|
|
"total_params": cc_params, |
|
|
"class_names": CLASS_NAMES, |
|
|
"text_encoder": QwenEmbeddingExtractor.MODEL_NAME, |
|
|
"voxel_encoder": "GeometricShapeClassifier_v8", |
|
|
} |
|
|
with open(cc_staging / "config.json", "w") as f: |
|
|
json.dump(cc_arch, f, indent=2) |
|
|
|
|
|
cc_train_cfg = { |
|
|
"n_samples": CC_SAMPLES, "epochs": CC_EPOCHS, "batch_size": CC_BATCH, |
|
|
"lr": CC_LR, "weight_decay": 1e-4, "optimizer": "AdamW", |
|
|
"scheduler": "cosine_with_warmup", "warmup_epochs": _warmup, |
|
|
"loss": "symmetric_InfoNCE", "initial_temperature": 0.07, |
|
|
"final_temperature": cc_model.temperature.item(), |
|
|
"amp_dtype": str(amp_dtype), |
|
|
"best_val_accuracy": best_cc_acc, |
|
|
"nearest_text_accuracy": correct / max(total_c, 1), |
|
|
"total_training_time_seconds": tt, |
|
|
} |
|
|
with open(cc_staging / "training_config.json", "w") as f: |
|
|
json.dump(cc_train_cfg, f, indent=2) |
|
|
|
|
|
|
|
|
upload_cc_to_hf(cc_model, best_cc_acc, CC_EPOCHS - 1, hf_token, reason="final", |
|
|
text_dim_=text_dim, voxel_dim_=voxel_dim, latent_dim_=CC_LATENT) |
|
|
if not hf_token: |
|
|
print("No HF_TOKEN — saved locally at ./hf_staging/crosscontrast/") |
|
|
|
|
|
print(f"\nAll three subdirectories staged:") |
|
|
print(f" geometric_classifier/ — from Cell 3") |
|
|
print(f" qwen_embeddings/ — {text_embeddings.shape}") |
|
|
print(f" crosscontrast/ — latent={CC_LATENT}, acc={best_cc_acc:.4f}") |
|
|
print(f" Repo: https://huggingface.co/{HF_REPO}") |
