depth_fm.training

training

Training loop and Lightning module for DepthFM.

• lightning_module.py — DepthFMLightningModule (train/val/test steps, EMA, checkpoint strategy).
• train_lightning.py — CLI entry point. Invoked via torchrun -m depth_fm.training.train_lightning.

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")