Upload portable training script

train_portable.py

@@ -0,0 +1,434 @@
+#!/usr/bin/env python3
+"""
+Portable DPSS Training Script.
+
+Runs on ANY GPU machine: local workstation, Colab, RunPod, Lambda, Vast.ai, etc.
+
+Workflow:
+    1. Downloads the ArchCAD-40K dataset from HuggingFace Hub
+    2. Downloads upstream DPSS code bundle from mohansshf/dpss-archcad
+    3. Downloads HRNet-W48 pretrained backbone weights
+    4. Trains DPSS model (PointTransformerV2 + HRNet48 + Mask2Former decoder)
+    5. Pushes best checkpoint to mohansshf/dpss-archcad (or saves locally)
+
+Prerequisites:
+    pip install torch torchvision numpy scipy Pillow munch pyyaml \\
+        tensorboardX tensorboard tqdm huggingface_hub
+    pip install 'git+https://github.com/facebookresearch/detectron2.git'
+    # Optional (speeds up pointops ~3x on GPU):
+    pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'
+
+Quick start:
+    python train_portable.py                       # Full training, batch_size=8
+    python train_portable.py --batch_size 4        # Smaller batch for 24GB GPUs
+    python train_portable.py --batch_size 2 --fp16 # For 16GB GPUs (T4/4090)
+    python train_portable.py --epochs 5 --no_push  # Quick test run
+"""
+
+import argparse
+import json as json_mod
+import os
+import re
+import signal
+import subprocess
+import sys
+import time
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Train DPSS model")
+    parser.add_argument("--batch_size", type=int, default=8,
+                        help="Per-GPU batch size (8 for 80GB, 4 for 24GB, 2 for 16GB)")
+    parser.add_argument("--epochs", type=int, default=100)
+    parser.add_argument("--patience", type=int, default=10,
+                        help="Early stopping patience (epochs without sPQ improvement)")
+    parser.add_argument("--fp16", action="store_true",
+                        help="Mixed precision training (helps on smaller GPUs)")
+    parser.add_argument("--lr", type=float, default=None,
+                        help="Learning rate (default: auto-scaled from config)")
+    parser.add_argument("--work_dir", type=str, default="/tmp/dpss_checkpoints",
+                        help="Directory for checkpoints")
+    parser.add_argument("--data_dir", type=str, default="/tmp/archcad_400k",
+                        help="Dataset directory (auto-downloaded if missing)")
+    parser.add_argument("--code_dir", type=str, default="/tmp/upstream",
+                        help="Upstream code directory (auto-downloaded if missing)")
+    parser.add_argument("--hub_repo", type=str, default="mohansshf/dpss-archcad",
+                        help="HF Hub repo to push weights to")
+    parser.add_argument("--no_push", action="store_true",
+                        help="Don't push to HF Hub (save locally only)")
+    parser.add_argument("--resume", type=str, default=None,
+                        help="Resume from a specific checkpoint path")
+    parser.add_argument("--install_deps", action="store_true",
+                        help="Install detectron2 and pytorch3d from source before training")
+    return parser.parse_args()
+
+
+def install_build_deps():
+    """Install detectron2 + pytorch3d from source."""
+    print("=" * 60)
+    print("INSTALLING BUILD DEPENDENCIES")
+    print("=" * 60)
+
+    print("\n>>> Installing detectron2 from source...")
+    subprocess.run(
+        [sys.executable, "-m", "pip", "install",
+         "git+https://github.com/facebookresearch/detectron2.git"],
+        check=True,
+    )
+
+    print("\n>>> Installing pytorch3d from source (may take 5-10 min)...")
+    try:
+        subprocess.run(
+            [sys.executable, "-m", "pip", "install",
+             "git+https://github.com/facebookresearch/pytorch3d.git@stable"],
+            check=True, timeout=900,
+        )
+        print("    pytorch3d installed (CUDA-accelerated pointops)")
+    except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
+        print("    WARNING: pytorch3d install failed — using pure-PyTorch pointops (slower but correct)")
+
+    print()
+
+
+def setup_data(data_dir):
+    """Download and extract ArchCAD dataset if not already present."""
+    json_dir = os.path.join(data_dir, "json")
+    png_dir = os.path.join(data_dir, "png")
+
+    if os.path.isdir(json_dir) and os.path.isdir(png_dir):
+        n_json = len([f for f in os.listdir(json_dir) if f.endswith(".json")])
+        n_png = len([f for f in os.listdir(png_dir) if f.endswith(".png")])
+        if n_json > 0 and n_png > 0:
+            print(f"Dataset already present: {n_json} json, {n_png} png")
+            return
+
+    from huggingface_hub import hf_hub_download
+    import zipfile
+
+    os.makedirs(data_dir, exist_ok=True)
+
+    for zipname in ["json.zip", "png.zip"]:
+        print(f">>> Downloading {zipname} from jackluoluo/ArchCAD...")
+        zip_path = hf_hub_download(
+            repo_id="jackluoluo/ArchCAD",
+            filename=f"data/{zipname}",
+            repo_type="dataset",
+        )
+        print(f"    Extracting -> {data_dir}")
+        with zipfile.ZipFile(zip_path, "r") as zf:
+            zf.extractall(data_dir)
+
+    n_json = len([f for f in os.listdir(json_dir) if f.endswith(".json")])
+    n_png = len([f for f in os.listdir(png_dir) if f.endswith(".png")])
+    print(f"    Dataset ready: {n_json} json, {n_png} png")
+    assert n_json > 0 and n_png > 0, "Dataset download/extraction failed!"
+
+
+def setup_code(code_dir):
+    """Download and extract upstream DPSS code if not already present."""
+    train_py = os.path.join(code_dir, "tools", "train.py")
+    if os.path.exists(train_py):
+        print(f"Upstream code already present at {code_dir}")
+        return
+
+    from huggingface_hub import hf_hub_download
+    import tarfile
+
+    os.makedirs(code_dir, exist_ok=True)
+
+    print(">>> Downloading upstream DPSS code from mohansshf/dpss-archcad...")
+    code_tar = hf_hub_download(
+        repo_id="mohansshf/dpss-archcad",
+        filename="upstream_code.tar.gz",
+        repo_type="model",
+    )
+    print(f"    Extracting -> {code_dir}")
+    with tarfile.open(code_tar, "r:gz") as tar:
+        tar.extractall(code_dir)
+
+    assert os.path.exists(train_py), f"Code extraction failed — {train_py} not found!"
+    print("    Code ready")
+
+
+def setup_pretrained(code_dir):
+    """Download HRNet-W48 pretrained weights."""
+    pretrained_dir = os.path.join(code_dir, "pretrained")
+    os.makedirs(pretrained_dir, exist_ok=True)
+    hrnet_path = os.path.join(pretrained_dir, "hrnet_ocr_cocostuff_3965_torch04.pth")
+
+    if os.path.exists(hrnet_path) and os.path.getsize(hrnet_path) > 1e6:
+        print(f"HRNet weights already present: {os.path.getsize(hrnet_path) / 1e6:.0f} MB")
+        return hrnet_path
+
+    hrnet_url = "https://github.com/hsfzxjy/models.storage/releases/download/HRNet-OCR/hrnet_ocr_cocostuff_3965_torch04.pth"
+    print(">>> Downloading HRNet-W48 pretrained weights...")
+    subprocess.run(["wget", "-q", "--show-progress", "-O", hrnet_path, hrnet_url], check=True)
+    print(f"    Downloaded: {os.path.getsize(hrnet_path) / 1e6:.0f} MB")
+    return hrnet_path
+
+
+def generate_config(data_dir, code_dir, hrnet_path, batch_size, epochs, fp16, lr=None):
+    """Generate runtime YAML config from upstream template."""
+    import yaml
+
+    config_path = os.path.join(code_dir, "configs", "svg", "svg_pointT.yaml")
+    with open(config_path) as f:
+        cfg = yaml.safe_load(f)
+
+    # Data paths
+    cfg["data"]["train"]["data_root"] = data_dir
+    cfg["data"]["test"]["data_root"] = data_dir
+
+    # Hyperparams
+    cfg["dataloader"]["train"]["batch_size"] = batch_size
+    cfg["dataloader"]["test"]["batch_size"] = max(batch_size, 4)  # eval needs >=4 for BatchNorm
+    cfg["epochs"] = epochs
+    cfg["fp16"] = fp16
+
+    # Learning rate
+    if lr is not None:
+        cfg["optimizer"]["lr"] = lr
+
+    # Pretrained weight paths (absolute)
+    cfg["model"]["vision"]["pretrained"] = hrnet_path
+    cfg["model"]["vision"]["yaml"] = os.path.join(code_dir, "configs", "svg", "hrnet48.yaml")
+
+    # Save
+    runtime_path = os.path.join(code_dir, "runtime_config.yaml")
+    with open(runtime_path, "w") as f:
+        yaml.dump(cfg, f, default_flow_style=False)
+
+    print(f"\nConfig: batch_size={batch_size}, epochs={epochs}, fp16={fp16}")
+    base_lr = lr or cfg["optimizer"]["lr"]
+    print(f"  Base LR: {base_lr} (will be auto-scaled by train.py for effective batch size)")
+    return runtime_path
+
+
+def train(code_dir, config_path, work_dir, epochs, patience):
+    """Launch training and monitor for early stopping + push checkpoints."""
+    import torch
+
+    print("\n" + "=" * 60)
+    print("DPSS TRAINING")
+    print("=" * 60)
+    if torch.cuda.is_available():
+        print(f"GPU:        {torch.cuda.get_device_name(0)}")
+        print(f"VRAM:       {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
+    else:
+        print("WARNING: No CUDA GPU detected! Training will be extremely slow.")
+    print(f"Epochs:     {epochs}")
+    print(f"Patience:   {patience}")
+    print(f"Work dir:   {work_dir}")
+    print()
+
+    os.makedirs(work_dir, exist_ok=True)
+
+    # Resume detection
+    resume_arg = []
+    latest_pth = os.path.join(work_dir, "latest.pth")
+    if os.path.exists(latest_pth):
+        print(f"Resuming from: latest.pth")
+        resume_arg = ["--resume", latest_pth]
+    else:
+        existing = sorted(
+            [f for f in os.listdir(work_dir) if f.startswith("epoch_") and f.endswith(".pth")],
+            key=lambda x: int(x.split("_")[1].split(".")[0])
+        ) if os.path.isdir(work_dir) else []
+        if existing:
+            latest = os.path.join(work_dir, existing[-1])
+            print(f"Resuming from: {latest}")
+            resume_arg = ["--resume", latest]
+        else:
+            print("Starting fresh training")
+
+    # Launch training subprocess
+    os.chdir(code_dir)
+    env = os.environ.copy()
+    env["PYTHONPATH"] = code_dir
+    env["PYTHONUNBUFFERED"] = "1"
+
+    cmd = [
+        sys.executable, "-m", "torch.distributed.run",
+        "--nproc_per_node=1",
+        "--master_port=29500",
+        "tools/train.py",
+        config_path,
+        "--dist",
+        "--work_dir", work_dir,
+    ] + resume_arg
+
+    print(f"CMD: {' '.join(cmd)}")
+    print("=" * 60 + "\n")
+
+    process = subprocess.Popen(
+        cmd, env=env,
+        stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
+        text=True, bufsize=1,
+        start_new_session=True,
+    )
+
+    # Regex for log parsing
+    spq_pattern = re.compile(r"INFO\s+-\s+PQ / RQ / SQ\s*:\s*([\d.]+)")
+    loss_pattern = re.compile(r"\bloss:\s*([-+\d.eEinfaN]+)")
+
+    best_spq = 0.0
+    epochs_no_improve = 0
+    spq_history = []
+    nan_count = 0
+
+    # Restore early stopping state
+    es_file = os.path.join(work_dir, "early_stopping.json")
+    if os.path.exists(es_file):
+        with open(es_file) as f:
+            es = json_mod.load(f)
+        best_spq = es.get("best_spq", 0)
+        epochs_no_improve = es.get("epochs_without_improvement", 0)
+        spq_history = es.get("history", [])
+
+    def kill_proc():
+        try:
+            os.killpg(os.getpgid(process.pid), signal.SIGTERM)
+        except Exception:
+            try: process.terminate()
+            except: pass
+
+    for line in process.stdout:
+        print(line, end="", flush=True)
+
+        # NaN detection
+        m = loss_pattern.search(line)
+        if m:
+            try:
+                v = float(m.group(1))
+                if v != v or abs(v) == float('inf'):
+                    nan_count += 1
+                    if nan_count >= 50:
+                        print(f"\n>>> ABORT: {nan_count} NaN/Inf losses\n")
+                        kill_proc()
+                        break
+            except: pass
+
+        # Early stopping
+        m = spq_pattern.search(line)
+        if m:
+            val = float(m.group(1))
+            spq_history.append(val)
+            if val > best_spq:
+                best_spq = val
+                epochs_no_improve = 0
+                print(f"\n>>> New best sPQ = {val:.3f}\n")
+            else:
+                epochs_no_improve += 1
+                print(f"\n>>> sPQ={val:.3f} <= best={best_spq:.3f}  ({epochs_no_improve}/{patience})\n")
+
+            with open(es_file, "w") as f:
+                json_mod.dump({"best_spq": best_spq, "epochs_without_improvement": epochs_no_improve, "history": spq_history}, f)
+
+            if epochs_no_improve >= patience:
+                print(f"\n{'='*60}\nEARLY STOPPING after {patience} epochs without improvement\nBest sPQ: {best_spq:.3f}\n{'='*60}\n")
+                kill_proc()
+                try: process.wait(timeout=30)
+                except: pass
+                break
+
+    process.wait()
+    print(f"\nTraining finished. Best sPQ: {best_spq:.3f}")
+    print(f"History: {[f'{s:.3f}' for s in spq_history]}")
+    return best_spq
+
+
+def export_and_push(work_dir, hub_repo, no_push=False):
+    """Export clean weights and optionally push to HF Hub."""
+    import torch
+
+    best_path = os.path.join(work_dir, "best.pth")
+    if not os.path.exists(best_path):
+        ckpts = sorted(
+            [f for f in os.listdir(work_dir) if f.startswith("epoch_") and f.endswith(".pth")],
+            key=lambda x: int(x.split("_")[1].split(".")[0])
+        )
+        if ckpts:
+            best_path = os.path.join(work_dir, ckpts[-1])
+        else:
+            print("ERROR: No checkpoints found!")
+            return None
+
+    print(f"\n>>> Exporting weights from {best_path}")
+    ckpt = torch.load(best_path, map_location="cpu")
+    weights = ckpt.get("net", ckpt.get("state_dict", ckpt.get("model", ckpt)))
+    epoch = ckpt.get("epoch", "unknown")
+
+    export_path = os.path.join(work_dir, "dpss_weights.pth")
+    torch.save({
+        "state_dict": weights,
+        "meta": {"source": "dpss-training", "dataset": "ArchCAD-40K", "epoch": epoch},
+    }, export_path)
+
+    n_params = sum(p.numel() for p in weights.values()) / 1e6
+    size_mb = os.path.getsize(export_path) / 1e6
+    print(f"    {n_params:.1f}M params, {size_mb:.1f} MB, epoch {epoch}")
+    print(f"    Saved to: {export_path}")
+
+    if not no_push:
+        try:
+            from huggingface_hub import HfApi
+            api = HfApi()
+            print(f"\n>>> Pushing to {hub_repo}...")
+            api.upload_file(path_or_fileobj=export_path, path_in_repo="dpss_weights.pth",
+                            repo_id=hub_repo, commit_message=f"DPSS weights (epoch {epoch})")
+            api.upload_file(path_or_fileobj=best_path, path_in_repo="best_full.pth",
+                            repo_id=hub_repo, commit_message=f"Full checkpoint (epoch {epoch})")
+            es_file = os.path.join(work_dir, "early_stopping.json")
+            if os.path.exists(es_file):
+                api.upload_file(path_or_fileobj=es_file, path_in_repo="training_metrics.json",
+                                repo_id=hub_repo, commit_message="Training metrics")
+            print(f"    Pushed to https://huggingface.co/{hub_repo}")
+        except Exception as e:
+            print(f"    WARNING: Push failed ({e}). Weights saved locally at {export_path}")
+
+    return export_path
+
+
+def main():
+    args = parse_args()
+    t0 = time.time()
+
+    print("=" * 60)
+    print("DPSS PORTABLE TRAINING")
+    print("=" * 60)
+
+    # Optional: install build deps
+    if args.install_deps:
+        install_build_deps()
+
+    # Setup
+    print("\n>>> Setting up data, code, and pretrained weights...\n")
+    setup_data(args.data_dir)
+    setup_code(args.code_dir)
+    hrnet_path = setup_pretrained(args.code_dir)
+
+    # Config
+    config_path = generate_config(
+        args.data_dir, args.code_dir, hrnet_path,
+        args.batch_size, args.epochs, args.fp16, args.lr,
+    )
+
+    # Train
+    best_spq = train(args.code_dir, config_path, args.work_dir, args.epochs, args.patience)
+
+    # Export
+    export_path = export_and_push(args.work_dir, args.hub_repo, args.no_push)
+
+    elapsed = time.time() - t0
+    print(f"\n{'='*60}")
+    print(f"COMPLETE — {elapsed/3600:.1f}h elapsed")
+    print(f"Best sPQ: {best_spq:.3f}")
+    if export_path:
+        print(f"Weights: {export_path}")
+    if not args.no_push:
+        print(f"Hub: https://huggingface.co/{args.hub_repo}")
+    print(f"{'='*60}")
+
+
+if __name__ == "__main__":
+    main()