depth_fm.training.lightning_module

lightning_module

PyTorch Lightning module for Mars DepthFM.

[REFACTOR NOTE] 1982 LOC. Candidate for future split into: lightning_module.py (train/val/test steps), training_metrics.py (aggregation), and visualization_worker.py (multiprocessing figure rendering). Out of scope for the structural refactor that produced this layout.

Implements the DepthFM flow matching training loop with: - Proper train/val/test step separation - Per-epoch metric logging with all standard depth metrics - Flow evolution visualisation at configurable intervals - Normals loss warmup scheduling

DDP-safe: all sync_dist logging is unconditional, all model inference in visualization runs on all ranks (only rank 0 logs figures).

──────────────────────────────────────────────────────────────────────── Visualization throughput design (v2 — multiprocessing + round-robin) ────────────────────────────────────────────────────────────────────────

WHY THE ORIGINAL WAS SLOW ────────────────────────── The original code used a concurrent.futures.ThreadPoolExecutor on rank 0 for background figure generation. Three problems:

1. GIL contention. matplotlib (Agg backend) is pure-Python / Cython and holds the GIL during every rasterisation call. The training DataLoader pre-fetch, metric aggregation, and gradient logging all compete for the same GIL on rank 0.

2. Rank asymmetry → DDP stalls. Only rank 0 generated figures, so it fell behind. Every subsequent DDP collective (all-reduce during backward, sync_dist logging, barrier at checkpointing) forced the other GPUs to idle until rank 0 caught up.

3. Blocking checkpoint flush. on_save_checkpoint called _flush_vis_tasks(timeout=120s), stalling the entire DDP group because the other ranks had already passed the save barrier and were waiting at the next collective.

HOW THIS VERSION FIXES IT ───────────────────────── A. torch.multiprocessing worker (one per DDP rank). A persistent daemon mp.Process (started with the spawn context so it never inherits CUDA state — PyTorch docs mandate this) runs a tight matplotlib→PDF/PNG loop. Because it is a process, not a thread, it owns its own GIL and never blocks the training process.

B. Round-robin figure distribution. Instead of funneling all 9+ figures through rank 0, we broadcast the tiny vis snapshot (~1-2 MB of numpy arrays) from rank 0 to all ranks via torch.distributed.broadcast, then assign figures to ranks with task_idx % world_size. Each rank's worker renders only its share of figures in parallel, cutting wall-clock time by ≈ world_size×. Single-GPU training degrades gracefully (all tasks stay on rank 0).

C. Deferred wandb upload. All ranks write PNGs + PDFs to a shared staging directory. At the start of the next validation epoch, rank 0 sweeps the staging directory and does a single batched wandb.log() call. This guarantees that uploads never block the training loop — by the time we look, the previous epoch's figures are certainly on disk.

D. Non-blocking checkpoint. on_save_checkpoint no longer flushes pending vis tasks. Worst case, a crash loses some PDFs from the current epoch — acceptable because the model weights are safe.

E. mp.SimpleQueue instead of mp.Queue. Per PyTorch multiprocessing docs, SimpleQueue spawns no background threads and avoids the deadlock-prone serialisation threads that mp.Queue uses internally.

PERFORMANCE IMPACT (expected) ───────────────────────────── • GPU 0 idle time during validation: eliminated (no longer generates figures synchronously or via GIL-holding threads). • Total figure generation wall time: ÷ world_size (round-robin). • wandb upload latency: moved entirely out of the training critical path (deferred to next epoch start). • Checkpoint save time: reduced by up to 120 s (no vis flush).

DepthFMLightningModule

DepthFMLightningModule(config)

Bases: LightningModule

Lightning module for DepthFM flow matching training.

Handles: - Conditional flow matching velocity loss - Surface normals auxiliary loss with warmup - Logit-normal timestep sampling - Noise augmentation on source distribution - Multi-step Euler inference for validation - Full metric computation and logging

Source code in src/depth_fm/training/lightning_module.py

def __init__(self, config):
    super().__init__()
    self.save_hyperparameters(config)
    self.config = config

    # Build model
    torch.backends.cudnn.benchmark = True
    self.model = build_model(config)

    if config.training.get("optimizer", "adamw") != "shampoo":
        self.model = self.model.to(memory_format=torch.channels_last)
    else:
        logger.warning("Disabled channels_last memory format for shampoo training")

    # Loss
    lc = config.training.losses
    self.loss_fn = CombinedLoss(
        velocity_weight=lc.velocity_weight,
        normals_weight=lc.get("normals_weight", 0.1),
        normals_start_step=lc.get("normals_start_step", 2000),
        use_confidence_weighting=lc.get("use_confidence_weighting", False),
        freq_weight=lc.get("freq_weight", 0.0),
        freq_start_step=lc.get("freq_start_step", 0),
        freq_alpha=lc.get("freq_alpha", 1.0),
        grad_weight=lc.get("grad_weight", 0.0),
        grad_start_step=lc.get("grad_start_step", 0),
        grad_scales=tuple(lc.get("grad_scales", [1, 2, 4])),
        photo_weight=lc.get("photo_weight", 0.0),
        photo_start_step=lc.get("photo_start_step", 2000),
        huber_weight=lc.get("huber_weight", 0.0),
        huber_start_step=lc.get("huber_start_step", 0.0),
        laplacian_weight=lc.get("laplacian_weight", 0.0),
        laplacian_start_step=lc.get("laplacian_start_step", 0.0),
        ordinal_weight=lc.get("ordinal_weight", 0.0),
        ordinal_start_step=lc.get("ordinal_start_step", 0.0),

    )

    # Elevation normaliser — used as fallback when batch has no residual_scale.
    # clip must match the adapter so clip=False data (values outside [-1,1]) is
    # denormalized correctly (clip has no effect on denormalize_prediction but
    # keeps the object consistent with the data that was produced).
    self.evel_normalizer = GlobalLogNormalizer(
        config.data.get("elev_ref_scale", DEFAULT_ELEV_REF_SCALE),
        clip=config.data.get("clip", False),
    )

    # Flow matching config
    self.fm = config.training.flow_matching

    # Metric aggregators (reset each epoch)
    self._val_aggregator = MetricsAggregator()
    self._test_aggregator = MetricsAggregator()

    # Optional tag suffix for test-output files (used by the ablation
    # orchestrator when running test twice — once per checkpoint type).
    # Empty string → test_per_patch.npz / test_results.json (default).
    self._test_tag: str = ""

    # Track validation history for convergence plotting
    self.val_history: dict[str, list[float]] = {
        "val/rmse": [], "val/abs_rel": [], "val/delta_1": [],
        "val/loss": [], "train/loss": [],
    }

    # Cache for grad norm logging (avoid recomputing every step)
    self._grad_norm_log_interval = 50

    # ── Visualisation system (initialised lazily in on_train_start) ──
    # These are set up once training begins and the Trainer context is
    # available (world_size, default_root_dir, logger, etc.).
    self._vis_worker: mp.Process | None = None
    self._vis_queue: mp.SimpleQueue | None = None
    self._vis_staging_dir: str | None = None
    self._vector_fig_dir: str | None = None

estimate_uncertainty

estimate_uncertainty(z_img: Tensor, num_samples: int = 10, num_steps: int = 4) -> tuple[torch.Tensor, torch.Tensor]

Computes the epistemic uncertainty via stochastic ODE sampling.

Parameters:

Name	Type	Description	Default
z_img	Tensor	(B, C, H, W) Encoded input orthoimage.	required
num_samples	int	Number of stochastic noise initializations (N).	10
num_steps	int	Euler integration steps per sample.	4

Returns:

Name	Type	Description
mean_dtm	Tensor	(B, 1, H, W) The expected topographic surface.
var_dtm	Tensor	(B, 1, H, W) The pixel-wise epistemic variance.

Source code in src/depth_fm/training/lightning_module.py

@torch.no_grad()
def estimate_uncertainty(
        self,
        z_img: torch.Tensor,
        num_samples: int = 10,
        num_steps: int = 4
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Computes the epistemic uncertainty via stochastic ODE sampling.

    Args:
        z_img: (B, C, H, W) Encoded input orthoimage.
        num_samples: Number of stochastic noise initializations (N).
        num_steps: Euler integration steps per sample.

    Returns:
        mean_dtm: (B, 1, H, W) The expected topographic surface.
        var_dtm: (B, 1, H, W) The pixel-wise epistemic variance.
    """
    self._trace(f"Starting epistemic uncertainty estimation with N={num_samples}")

    dtm_hypotheses = []

    for i in range(num_samples):
        # _predict_depth automatically samples a new x_source ~ N(0, I)
        # via the _get_x_source() method internally.
        z_pred = self._predict_depth(z_img, num_steps=num_steps)

        pred_pix = self.evel_normalizer.denormalize_prediction(
            self._decode(z_pred)[:, 0])
        dtm_hypotheses.append(pred_pix)

    # Stack into shape: (N, B, 1, H, W)
    dtm_tensor = torch.stack(dtm_hypotheses, dim=0)

    # Calculate moments
    mean_dtm = torch.mean(dtm_tensor, dim=0)
    var_dtm = torch.var(dtm_tensor, dim=0, unbiased=True)

    return mean_dtm, var_dtm

on_train_end

on_train_end() -> None

Final wandb upload and vis worker shutdown.

Source code in src/depth_fm/training/lightning_module.py

def on_train_end(self) -> None:
    """Final wandb upload and vis worker shutdown."""
    # Upload any remaining staged figures
    self._deferred_wandb_upload()
    self._upload_vector_figures_artifact(prefix="val")

    # Shut down the worker process
    self._shutdown_vis_worker(timeout=60.0)
    self._trace("Training complete, vis worker shut down")