Upload learn_region_grow/train.py

learn_region_grow/train.py

@@ -0,0 +1,175 @@
+"""Training script for LrgNet."""
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import Dataset, DataLoader
+import numpy as np
+import h5py
+from pathlib import Path
+from typing import List, Optional
+from .lrg_net import LrgNet
+
+
+class StagedDataset(Dataset):
+    """PyTorch Dataset wrapping H5 staged training files."""
+
+    def __init__(self, h5_paths: List[str]):
+        self.h5_paths = h5_paths
+        self.offsets = []
+        self.total = 0
+        for p in h5_paths:
+            with h5py.File(p, 'r') as f:
+                n = f['inliers'].shape[0]
+                self.offsets.append((self.total, self.total + n, p))
+                self.total += n
+
+    def __len__(self):
+        return self.total
+
+    def __getitem__(self, idx):
+        # Find which file
+        for start, end, path in self.offsets:
+            if start <= idx < end:
+                local_idx = idx - start
+                break
+        else:
+            raise IndexError(idx)
+
+        with h5py.File(path, 'r') as f:
+            inlier = f['inliers'][local_idx]      # (Ni, 13)
+            neighbor = f['neighbors'][local_idx]  # (Nn, 13)
+            add_lbl = f['add_labels'][local_idx]    # (Nn,)
+            rmv_lbl = f['remove_labels'][local_idx] # (Ni,)
+
+        # Transpose to (C, N) for Conv1d
+        inlier = torch.from_numpy(inlier.T).float()    # (13, Ni)
+        neighbor = torch.from_numpy(neighbor.T).float()  # (13, Nn)
+        add_lbl = torch.from_numpy(add_lbl).long()       # (Nn,)
+        rmv_lbl = torch.from_numpy(rmv_lbl).long()       # (Ni,)
+
+        return inlier, neighbor, add_lbl, rmv_lbl
+
+
+class AddRemoveLoss(nn.Module):
+    """Joint cross-entropy over add + remove logits."""
+
+    def __init__(self, add_weight: float = 1.0, remove_weight: float = 1.0):
+        super().__init__()
+        self.add_weight = add_weight
+        self.remove_weight = remove_weight
+        self.ce = nn.CrossEntropyLoss(reduction='none')
+
+    def forward(self, add_logits, add_targets, remove_logits, remove_targets):
+        # add_logits: (B, 1, Nn) -> treat as binary classification
+        # PyTorch cross_entropy expects (B, C, ...); here C=1 is tricky for sigmoid
+        # Simpler: use BCEWithLogitsLoss
+        pass  # placeholder -- we use BCE in trainer below
+
+
+def train_lrgnet(train_files: List[str],
+                 val_files: Optional[List[str]] = None,
+                 epochs: int = 50,
+                 batch_size: int = 16,
+                 lr: float = 1e-3,
+                 device: str = 'cuda',
+                 lite: int = 0,
+                 save_dir: str = './checkpoints',
+                 resume: Optional[str] = None):
+    """
+    Train LrgNet on staged H5 files.
+
+    Parameters
+    ----------
+    train_files : list of str
+        Paths to staged H5 files.
+    val_files : list of str, optional
+        Validation H5 files.
+    epochs : int
+    batch_size : int
+    lr : float
+        Adam learning rate (default 1e-3, matching the paper).
+    device : str
+    lite : int
+        0 = full, 1 = half channels, 2 = quarter channels.
+    save_dir : str
+        Where to write checkpoints.
+    resume : str, optional
+        Path to a checkpoint to resume from.
+    """
+    device = torch.device(device if torch.cuda.is_available() else 'cpu')
+
+    train_ds = StagedDataset(train_files)
+    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,
+                              num_workers=4, pin_memory=True, drop_last=True)
+
+    val_loader = None
+    if val_files:
+        val_ds = StagedDataset(val_files)
+        val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False,
+                                num_workers=4, pin_memory=True)
+
+    model = LrgNet(in_channels=13, lite=lite).to(device)
+    if resume:
+        model.load_state_dict(torch.load(resume, map_location=device))
+
+    optimizer = optim.Adam(model.parameters(), lr=lr)
+    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.5)
+
+    bce_add = nn.BCEWithLogitsLoss()
+    bce_remove = nn.BCEWithLogitsLoss()
+
+    save_dir = Path(save_dir)
+    save_dir.mkdir(parents=True, exist_ok=True)
+
+    best_val_loss = float('inf')
+
+    for epoch in range(1, epochs + 1):
+        model.train()
+        total_loss = 0.0
+        n_batches = 0
+
+        for inliers, neighbors, add_lbl, rmv_lbl in train_loader:
+            inliers = inliers.to(device)
+            neighbors = neighbors.to(device)
+            add_lbl = add_lbl.to(device).float().unsqueeze(1)  # (B, 1, Nn)
+            rmv_lbl = rmv_lbl.to(device).float().unsqueeze(1)  # (B, 1, Ni)
+
+            optimizer.zero_grad()
+            add_logits, rmv_logits = model(inliers, neighbors)
+
+            # For binary BCE, logits shape is (B, 1, N). Targets same shape.
+            loss = bce_add(add_logits, add_lbl) + bce_remove(rmv_logits, rmv_lbl)
+            loss.backward()
+            optimizer.step()
+
+            total_loss += loss.item()
+            n_batches += 1
+
+        avg_train = total_loss / max(n_batches, 1)
+
+        val_str = ""
+        if val_loader:
+            model.eval()
+            val_loss = 0.0
+            with torch.no_grad():
+                for inliers, neighbors, add_lbl, rmv_lbl in val_loader:
+                    inliers = inliers.to(device)
+                    neighbors = neighbors.to(device)
+                    add_lbl = add_lbl.to(device).float().unsqueeze(1)
+                    rmv_lbl = rmv_lbl.to(device).float().unsqueeze(1)
+                    add_logits, rmv_logits = model(inliers, neighbors)
+                    vloss = bce_add(add_logits, add_lbl) + bce_remove(rmv_logits, rmv_lbl)
+                    val_loss += vloss.item()
+            avg_val = val_loss / len(val_loader)
+            val_str = f" | val_loss={avg_val:.4f}"
+            if avg_val < best_val_loss:
+                best_val_loss = avg_val
+                torch.save(model.state_dict(), save_dir / 'best_model.pt')
+
+        scheduler.step()
+        print(f"Epoch {epoch}/{epochs}  train_loss={avg_train:.4f}{val_str}")
+        torch.save(model.state_dict(), save_dir / f'epoch_{epoch:03d}.pt')
+
+    print(f"Training complete. Checkpoints saved to {save_dir}")
+    return model