Faz2 dep: code/meshai_train/engine.py

code/meshai_train/engine.py

@@ -0,0 +1,188 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader, Subset
+
+from meshai_train.dataset import PreprocessedMeshDataset, collate_preprocessed
+from meshai_train.models import MeshAIHybridTrainBundle
+
+TRAIN_VERSION = "v4.0-real-preprocessed"
+
+
+def hybrid_loss(out: dict[str, torch.Tensor], batch: dict[str, torch.Tensor]) -> tuple[torch.Tensor, dict[str, float]]:
+    voxel_loss = nn.functional.mse_loss(out["voxel_pred"], batch["voxel_tgt"])
+    bridge_loss = nn.functional.mse_loss(out["bridge_out"], out["tex_latent"].detach())
+    tex_reg = out["tex_latent"].pow(2).mean() * 1e-4
+    total = voxel_loss + bridge_loss + tex_reg
+    return total, {
+        "voxel": float(voxel_loss.item()),
+        "bridge": float(bridge_loss.item()),
+        "tex_reg": float(tex_reg.item()),
+    }
+
+
+def save_real_checkpoint(
+    path: Path,
+    *,
+    epoch: int,
+    global_step: int,
+    model: MeshAIHybridTrainBundle,
+    extra: dict[str, Any] | None = None,
+) -> None:
+    payload = {
+        "version": TRAIN_VERSION,
+        "epoch": epoch,
+        "global_step": global_step,
+        "geometry": model.geometry.state_dict(),
+        "texture": model.texture.state_dict(),
+        "bridge": model.bridge.state_dict(),
+        "extra": extra or {},
+    }
+    torch.save(payload, path)
+
+
+def load_real_checkpoint(path: Path, model: MeshAIHybridTrainBundle, device: str) -> int:
+    state = torch.load(path, map_location=device, weights_only=False)
+    if state.get("version") != TRAIN_VERSION:
+        return 0
+    model.geometry.load_state_dict(state["geometry"])
+    model.texture.load_state_dict(state["texture"])
+    model.bridge.load_state_dict(state["bridge"])
+    return int(state.get("global_step", 0))
+
+
+def start_real_training(
+    *,
+    monitor: Any,
+    checkpoint_dir: Path,
+    output_dir: Path,
+    token: str | None,
+    epochs: int,
+    resume_from: Path | None,
+    validation_every: int,
+    val_ratio: float,
+    limit: int | None,
+    hf_repo: str,
+    data_root: Path | None,
+    log_fn: Any,
+    log_vram_fn: Any,
+    clear_gpu_fn: Any,
+    load_val_uids_fn: Any,
+    checkpoint_every: int = 100,
+) -> None:
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    dtype = torch.float32
+
+    log_fn(f"Pipeline surumu: {TRAIN_VERSION} (gercek preprocessed latent + render)")
+    log_fn(f"Veri: {hf_repo}" + (f" | limit={limit}" if limit else ""))
+    if torch.cuda.is_available():
+        log_fn(f"GPU: {torch.cuda.get_device_name(0)}")
+    log_vram_fn("startup")
+
+    dataset = PreprocessedMeshDataset(
+        token=token,
+        data_root=data_root,
+        hf_repo=hf_repo,
+        limit=limit,
+    )
+    val_uids = load_val_uids_fn()
+    val_idx = [i for i, o in enumerate(dataset.objects) if str(o.get("uid")) in val_uids]
+    if not val_idx:
+        val_count = max(1, int(len(dataset) * val_ratio))
+        val_idx = list(range(val_count))
+    train_idx = [i for i in range(len(dataset)) if i not in set(val_idx)]
+    if not train_idx:
+        train_idx = list(range(len(dataset)))
+        val_idx = train_idx[:1]
+
+    train_loader = DataLoader(
+        Subset(dataset, train_idx),
+        batch_size=2 if device == "cuda" else 1,
+        shuffle=True,
+        pin_memory=device == "cuda",
+        collate_fn=collate_preprocessed,
+    )
+    val_loader = DataLoader(
+        Subset(dataset, val_idx),
+        batch_size=2 if device == "cuda" else 1,
+        shuffle=False,
+        pin_memory=device == "cuda",
+        collate_fn=collate_preprocessed,
+    )
+    log_fn(f"Veri: {len(train_idx)} train + {len(val_idx)} val (gercek latent/render).")
+
+    model = MeshAIHybridTrainBundle().to(device=device, dtype=dtype)
+    opt = torch.optim.AdamW(model.parameters(), lr=1e-4, fused=False)
+    latest = checkpoint_dir / "latest_model.pt"
+    global_step = 0
+    if resume_from and resume_from.exists():
+        global_step = load_real_checkpoint(resume_from, model, device)
+        log_fn(f"Resume: {resume_from} step={global_step}")
+
+    def _eval() -> tuple[float, float]:
+        model.eval()
+        losses: list[float] = []
+        with torch.no_grad():
+            for batch in val_loader:
+                batch = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}
+                out = model(batch["geom_in"], batch["views"])
+                loss, _ = hybrid_loss(out, batch)
+                losses.append(float(loss.item()))
+        model.train()
+        mean = float(np.mean(losses)) if losses else float("nan")
+        return mean, mean
+
+    for epoch in range(1, epochs + 1):
+        log_fn(f"--- Epoch {epoch}/{epochs} ---")
+        model.train()
+        epoch_losses: list[float] = []
+        for batch in train_loader:
+            batch = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}
+            opt.zero_grad(set_to_none=True)
+            out = model(batch["geom_in"], batch["views"])
+            loss, parts = hybrid_loss(out, batch)
+            if not torch.isfinite(loss):
+                monitor.note_nan_skip("hybrid")
+                continue
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            opt.step()
+            global_step += 1
+            epoch_losses.append(float(loss.item()))
+            monitor.note_step(global_step, "geometry", parts["voxel"])
+            monitor.last["train_loss_texture"] = round(parts["bridge"], 6)
+            if global_step <= 3 or global_step % 25 == 0:
+                log_fn(
+                    f"step={global_step} loss={float(loss.item()):.6f} "
+                    f"voxel={parts['voxel']:.6f} bridge={parts['bridge']:.6f}"
+                )
+
+            if checkpoint_every > 0 and global_step % checkpoint_every == 0:
+                save_real_checkpoint(latest, epoch=epoch, global_step=global_step, model=model)
+                log_fn(
+                    f"CHECKPOINT_SAVED step={global_step} -> {latest} "
+                    f"({latest.stat().st_size // 1024} KB)"
+                )
+
+            if global_step % validation_every == 0:
+                val_geom, val_tex = _eval()
+                monitor.note_validation(global_step, val_geom, val_tex)
+                log_vram_fn(f"step_{global_step}")
+
+        val_geom, val_tex = _eval()
+        geom_mean = float(np.mean(epoch_losses)) if epoch_losses else float("nan")
+        monitor.note_epoch_end(epoch, epochs, geom_mean, val_tex, val_geom, val_tex)
+        save_real_checkpoint(latest, epoch=epoch, global_step=global_step, model=model)
+        clear_gpu_fn()
+        log_fn(f"Epoch {epoch} kaydedildi -> {latest} ({latest.stat().st_size // 1024} KB)")
+        log_fn(f"CHECKPOINT_SAVED step={global_step} epoch={epoch} -> {latest}")
+
+    monitor.finish(ok=True)
+    log_fn("Gercek preprocessed egitim tamamlandi.")
+    log_fn("Not: Tam TRELLIS/Hunyuan agirlik fine-tune sonraki adim; bu asama MeshAI hibrit latent koprusu.")