| |
| |
| import json |
| import random |
| import time |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from torch.utils.data import DataLoader |
|
|
| from .data import CachedDataset, image_tensor, load_manifest |
| from .model import Student, distillation_loss |
| from .splats import Splats, write_ply |
|
|
|
|
| def device_for(name="auto"): |
| if name == "auto": |
| name = "mps" if torch.backends.mps.is_available() else "cpu" |
| if name == "mps" and not torch.backends.mps.is_available(): |
| raise ValueError("Metal GPU unavailable in this process. Run in a normal macOS Terminal, or use --device cpu.") |
| if name not in ("mps", "cpu"): |
| raise ValueError("This prototype supports cpu and mps") |
| torch.set_num_threads(min(4, torch.get_num_threads())) |
| return torch.device(name) |
|
|
|
|
| def split_rows(rows, fraction, seed): |
| if not 0 <= fraction < 1: |
| raise ValueError("Validation fraction must lie in [0,1)") |
| |
| groups = {} |
| for row in rows: |
| group = row.get("metadata", {}).get("object_id", row["id"]) |
| groups.setdefault(group, []).append(row) |
| keys = sorted(groups) |
| random.Random(seed).shuffle(keys) |
| n = max(1, round(len(keys) * fraction)) if fraction else 0 |
| if n >= len(keys): |
| raise ValueError("Need at least two different objects for validation, or explicitly use --validation-fraction 0 for an overfit test") |
| val_keys = set(keys[:n]) |
| train = [r for k in keys if k not in val_keys for r in groups[k]] |
| val = [r for k in keys if k in val_keys for r in groups[k]] |
| if {r["image_sha256"] for r in train} & {r["image_sha256"] for r in val}: |
| raise ValueError("Identical images occur across training/validation. Assign the same object-id to related samples.") |
| return train, val |
|
|
|
|
| def on_device(target, device): |
| return {k: v.to(device) for k, v in target.items()} |
|
|
|
|
| @torch.no_grad() |
| def evaluate(model, loader, device): |
| model.eval() |
| sums = {} |
| for image, target in loader: |
| target = on_device(target, device) |
| pred = model(image.to(device)) |
| loss, terms = distillation_loss(pred, target, model.grid) |
| occupied = pred["occ"].sigmoid() > .5 |
| truth = target["occ"] > .5 |
| terms = {**terms, "loss": loss, "occupancy_iou": (occupied & truth).sum() / (occupied | truth).sum().clamp_min(1)} |
| for key, value in terms.items(): |
| sums[key] = sums.get(key, 0.) + float(value) |
| return {k: v / len(loader) for k, v in sums.items()} |
|
|
|
|
| def save_checkpoint(path, model, optimizer, step, metadata): |
| payload = {"format": 1, "config": model.config, "model": {k: v.detach().cpu() for k, v in model.state_dict().items()}, |
| "optimizer": optimizer.state_dict(), "step": step, "metadata": metadata} |
| tmp = Path(path).with_suffix(".tmp") |
| torch.save(payload, tmp) |
| tmp.replace(path) |
|
|
|
|
| def train(args): |
| if args.steps < 1 or args.batch_size < 1 or args.lr <= 0 or args.eval_every < 1: |
| raise ValueError("Steps, batch size, learning rate, and evaluation interval must be positive") |
| device = device_for(args.device) |
| random.seed(args.seed) |
| np.random.seed(args.seed) |
| torch.manual_seed(args.seed) |
| rows = load_manifest(args.manifest) |
| synthetic = any(r["source"] == "synthetic" for r in rows) |
| if synthetic and not args.allow_synthetic: |
| raise ValueError("Synthetic fixtures are code tests, not teacher data. Add --allow-synthetic for an explicit smoke test.") |
| train_rows, val_rows = split_rows(rows, args.validation_fraction, args.seed) |
| train_data = CachedDataset(args.manifest, train_rows) |
| generator = torch.Generator().manual_seed(args.seed) |
| loader = DataLoader(train_data, batch_size=args.batch_size, shuffle=True, num_workers=0, generator=generator) |
| val_loader = DataLoader(CachedDataset(args.manifest, val_rows), batch_size=1) if val_rows else None |
| model = Student(train_data.grid, args.width).to(device) |
| optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4) |
| run = Path(args.out) |
| if run.exists() and any(run.iterdir()): |
| raise ValueError("Run directory must be empty; choose a new --out directory") |
| run.mkdir(parents=True, exist_ok=True) |
| meta = {"synthetic_only": all(r["source"] == "synthetic" for r in rows), |
| "contains_synthetic": synthetic, "device": str(device), "torch": str(torch.__version__), |
| "seed": args.seed, "train_ids": [r["id"] for r in train_rows], "val_ids": [r["id"] for r in val_rows], |
| "parameters": sum(p.numel() for p in model.parameters()), |
| "teachers": sorted({v for r in rows for v in r["teachers"].values() if v}), |
| "validation": "held-out objects" if val_rows else "none; training overfit test only", |
| "dataset": str(Path(args.manifest).resolve())} |
| (run / "run.json").write_text(json.dumps(meta, indent=2)) |
| print(json.dumps(meta), flush=True) |
| step, best, started = 0, float("inf"), time.monotonic() |
| iterator = iter(loader) |
| try: |
| while step < args.steps: |
| try: |
| images, target = next(iterator) |
| except StopIteration: |
| iterator = iter(loader) |
| images, target = next(iterator) |
| model.train() |
| optimizer.zero_grad(set_to_none=True) |
| pred = model(images.to(device)) |
| loss, terms = distillation_loss(pred, on_device(target, device), model.grid) |
| if not torch.isfinite(loss): |
| raise RuntimeError("Nonfinite training loss; stopping") |
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(model.parameters(), 1.) |
| optimizer.step() |
| step += 1 |
| if step == 1 or step % args.eval_every == 0 or step == args.steps: |
| record = {"step": step, "seconds": round(time.monotonic()-started, 3), |
| "train_loss": float(loss.detach()), **{k: float(v.detach()) for k, v in terms.items()}} |
| if val_loader: |
| record["validation"] = evaluate(model, val_loader, device) |
| score = record["validation"]["loss"] |
| if score < best: |
| best = score |
| save_checkpoint(run / "best.pt", model, optimizer, step, meta) |
| with (run / "metrics.jsonl").open("a") as handle: |
| handle.write(json.dumps(record) + "\n") |
| print(json.dumps(record), flush=True) |
| save_checkpoint(run / "last.pt", model, optimizer, step, meta) |
| except KeyboardInterrupt: |
| print("Interrupted; saving completed optimizer steps.", flush=True) |
| finally: |
| if step: |
| save_checkpoint(run / "last.pt", model, optimizer, step, meta) |
| return meta |
|
|
|
|
| @torch.no_grad() |
| def generate(checkpoint, image, out, device="auto", threshold=.5, max_splats=32768): |
| if not 0 < threshold < 1 or max_splats < 1: |
| raise ValueError("Threshold must be in (0,1) and max-splats positive") |
| device = device_for(device) |
| state = torch.load(checkpoint, map_location="cpu", weights_only=True) |
| if state.get("format") != 1 or state.get("step", 0) < 1: |
| raise ValueError("Expected a trained image3d-student checkpoint") |
| model = Student(**state["config"]) |
| model.load_state_dict(state["model"]) |
| model.to(device).eval() |
| if device.type == "mps": |
| torch.mps.synchronize() |
| started = time.monotonic() |
| pred = model(image_tensor(image)[None].to(device)) |
| score = pred["occ"][0, :, 0].sigmoid() |
| valid = torch.nonzero(score >= threshold).flatten() |
| if len(valid) == 0: |
| raise ValueError("No occupied cells above threshold. Train longer or inspect validation; no fabricated fallback splats were exported.") |
| valid = valid[torch.argsort(score[valid], descending=True)[:max_splats]] |
| values = {k: pred[k][0, valid].cpu().numpy() for k in ("xyz", "rgb", "scale", "quat", "opacity")} |
| elapsed = time.monotonic() - started |
| splats = Splats(**values) |
| out = Path(out) |
| if out.exists() or out.with_suffix(".json").exists(): |
| raise ValueError("Output already exists; choose a new filename") |
| write_ply(out, splats) |
| report = {"gaussians": len(valid), "seconds_preprocess_and_forward": elapsed, |
| "device": str(device), "checkpoint_step": state["step"], "training": state["metadata"], |
| "note": ("Synthetic-trained output is a pipeline demonstration, not a pretrained image-to-3D result." |
| if state["metadata"].get("contains_synthetic") else "Experimental cached-target student; inspect unseen-object quality.")} |
| out.with_suffix(".json").write_text(json.dumps(report, indent=2)) |
| return report |
|
|
