"""Reproducible training and evaluation engine for scene-state experiments.""" from __future__ import annotations import hashlib import json import math import os import random import time from collections.abc import Iterator, Mapping from dataclasses import asdict from pathlib import Path import numpy as np import torch from torch import Tensor from torch.optim import AdamW from torch.utils.data import DataLoader, Dataset from mcss.config import DatasetConfig, ExperimentConfig from mcss.context_subset_consistency import ( compute_context_subset_geometry, drop_context_view, ramped_context_subset_geometry_weight, ) from mcss.data.collate import collate_scene_examples from mcss.data.manifest_dataset import ManifestSceneDataset from mcss.data.synthetic import SyntheticSceneDataset from mcss.losses import MeasurementLoss from mcss.metrics import compute_metrics from mcss.model.system import MeasurementCompleteSystem, build_model from mcss.types import SceneBatch, SceneExample, SceneState DEFAULT_LOSS_WEIGHTS = { "rgb": 1.0, "depth": 1.0, "normal": 0.5, "point": 0.25, "visibility": 0.5, "uncertainty": 0.1, } class _ScalarTracker: """Track log-window means and an EMA that can be restored from checkpoints.""" def __init__(self, *, ema_decay: float) -> None: if not 0.0 <= ema_decay < 1.0: raise ValueError("ema_decay must be in [0, 1)") self.ema_decay = float(ema_decay) self.window_sums: dict[str, float] = {} self.window_count = 0 self.ema: dict[str, float] = {} def update(self, values: Mapping[str, float]) -> None: self.window_count += 1 for name, raw_value in values.items(): value = float(raw_value) self.window_sums[name] = self.window_sums.get(name, 0.0) + value previous = self.ema.get(name) self.ema[name] = ( value if previous is None else self.ema_decay * previous + (1.0 - self.ema_decay) * value ) def snapshot(self, *, reset_window: bool) -> dict[str, float]: result = { **{ f"window/{name}": total / max(self.window_count, 1) for name, total in self.window_sums.items() }, **{f"ema/{name}": value for name, value in self.ema.items()}, } if reset_window: self.window_sums.clear() self.window_count = 0 return result def state_dict(self) -> dict[str, object]: return { "ema_decay": self.ema_decay, "window_sums": dict(self.window_sums), "window_count": self.window_count, "ema": dict(self.ema), } def load_state_dict(self, state: Mapping[str, object]) -> None: saved_decay = float(state.get("ema_decay", self.ema_decay)) if not math.isclose(saved_decay, self.ema_decay): raise ValueError("checkpoint telemetry ema_decay does not match config") self.window_sums = _float_mapping(state.get("window_sums", {}), "window_sums") self.window_count = int(state.get("window_count", 0)) self.ema = _float_mapping(state.get("ema", {}), "ema") class Trainer: """Own model, optimizer, data, output artifacts, and deterministic runtime state.""" def __init__(self, config: ExperimentConfig) -> None: self.config = config seed_everything(config.seed) self.device = resolve_device(config.device) self.output_dir = Path(config.training.output_dir) self.checkpoint_dir = self.output_dir / "checkpoints" self.model: MeasurementCompleteSystem = build_model(config.model).to(self.device) self.optimizer = AdamW( self.model.parameters(), lr=config.training.learning_rate, weight_decay=config.training.weight_decay, ) self.amp_enabled = config.training.amp and self.device.type == "cuda" self.amp_dtype = ( torch.bfloat16 if config.training.amp_dtype == "bfloat16" else torch.float16 ) if ( self.amp_enabled and self.amp_dtype == torch.bfloat16 and not torch.cuda.is_bf16_supported() ): raise RuntimeError( "bfloat16 AMP was requested but this CUDA device does not support it" ) self.scaler = torch.amp.GradScaler( "cuda", enabled=self.amp_enabled and self.amp_dtype == torch.float16 ) self.optimizer_step = 0 self.total_optimizer_steps = ( config.training.max_steps // config.training.gradient_accumulation ) self.telemetry = _ScalarTracker(ema_decay=config.training.ema_decay) self.context_subset_generator: torch.Generator | None = None if config.training.context_subset_geometry_weight > 0.0: self.context_subset_generator = torch.Generator(device="cpu") self.context_subset_generator.manual_seed(config.seed + 0x5C5C) weights = { name: DEFAULT_LOSS_WEIGHTS.get(name, 1.0) for name in config.training.measurements } self.objective = MeasurementLoss(weights) self.dataset = build_dataset(config.dataset, seed=config.seed) self.loader = DataLoader( self.dataset, batch_size=config.training.batch_size, shuffle=True, num_workers=config.training.num_workers, pin_memory=self.device.type == "cuda", collate_fn=collate_scene_examples, drop_last=False, ) self.step = 0 self.loaded_checkpoint_path: Path | None = None if config.training.resume: self.load_checkpoint(config.training.resume, load_optimizer=True) def train(self) -> Path: self.output_dir.mkdir(parents=True, exist_ok=True) self.checkpoint_dir.mkdir(parents=True, exist_ok=True) _atomic_json(self.output_dir / "resolved_config.json", asdict(self.config)) self.model.train() batches = _cycle(self.loader) batches_per_epoch = len(self.loader) self.optimizer.zero_grad(set_to_none=True) latest_checkpoint = self._checkpoint_path(self.step) for step in range(self.step + 1, self.config.training.max_steps + 1): step_started = time.perf_counter() batch = next(batches).to(self.device) with torch.autocast( device_type=self.device.type, dtype=self.amp_dtype, enabled=self.amp_enabled, ): output = self.model( batch.context_rgb, batch.context_cameras, batch.target_cameras, batch.bounds ) predictions = output.predictions targets = batch.targets() loss, terms = self.objective(predictions, targets) regularization = _state_regularization( output.state, evidence_residual_weight=self.config.training.evidence_residual_weight, ) objective_total = loss + regularization context_subset_result = None context_subset_metadata = None context_subset_weight = 0.0 if self.context_subset_generator is not None: dropped_index = int( torch.randint( batch.context_rgb.shape[1], (1,), generator=self.context_subset_generator, ).item() ) subset = drop_context_view( batch.context_rgb, batch.context_cameras, dropped_index=dropped_index, ) subset_state = self.model.state_encoder( subset.rgb, subset.cameras, batch.bounds, ) context_subset_result = compute_context_subset_geometry( output.state, subset_state, predictions, batch.target_cameras, self.model.renderer, residual_saturation_threshold=( 0.95 * self.config.model.completion_residual_scale ), collect_audit_diagnostics=( step % self.config.training.log_every == 0 ), ) context_subset_weight = ramped_context_subset_geometry_weight( self.config.training.context_subset_geometry_weight, optimizer_step=self.optimizer_step + 1, warmup_steps=self.config.training.warmup_steps, ) objective_total = ( objective_total + context_subset_weight * context_subset_result.loss ) context_subset_metadata = { "dropped_index": subset.dropped_index, "retained_indices": list(subset.retained_indices), "per_scene": [ { "scene_id": ( batch.scene_ids[index] if batch.scene_ids else None ), **diagnostics, } for index, diagnostics in enumerate( context_subset_result.per_sample_diagnostics ) ], } total = objective_total / self.config.training.gradient_accumulation self.scaler.scale(total).backward() optimizer_updated = False optimizer_skipped = False grad_norm: float | None = None if step % self.config.training.gradient_accumulation == 0: next_optimizer_step = self.optimizer_step + 1 learning_rate = self._learning_rate(next_optimizer_step) for parameter_group in self.optimizer.param_groups: parameter_group["lr"] = learning_rate self.scaler.unscale_(self.optimizer) norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0) grad_norm = float(norm.detach().item()) scale_before = float(self.scaler.get_scale()) self.scaler.step(self.optimizer) self.scaler.update() scale_after = float(self.scaler.get_scale()) optimizer_updated = not self.scaler.is_enabled() or scale_after >= scale_before optimizer_skipped = not optimizer_updated if optimizer_updated: self.optimizer_step = next_optimizer_step self.optimizer.zero_grad(set_to_none=True) self.step = step scalar_values = { "loss": float(loss.detach().item()), "regularization": float(regularization.detach().item()), "loss/total": float(objective_total.detach().item()), **{f"loss/{name}": float(value.detach().item()) for name, value in terms.items()}, **_state_diagnostics(output.state), "time/step_seconds": time.perf_counter() - step_started, } if context_subset_result is not None: scalar_values.update({ "loss/context_subset_geometry": float( context_subset_result.loss.detach().item() ), "loss/context_subset_depth": float( context_subset_result.terms["depth"].detach().item() ), "loss/context_subset_visibility": float( context_subset_result.terms["visibility"].detach().item() ), "context_subset/effective_weight": context_subset_weight, "context_subset/weighted_loss": float( (context_subset_weight * context_subset_result.loss).detach().item() ), **{ f"context_subset/{name}": value for name, value in context_subset_result.diagnostics.items() }, }) self.telemetry.update(scalar_values) if step % self.config.training.log_every == 0: epoch = (step - 1) // batches_per_epoch + 1 batch_in_epoch = (step - 1) % batches_per_epoch + 1 _append_jsonl( self.output_dir / "train.jsonl", { "step": step, "epoch": epoch, "epoch_progress": step / batches_per_epoch, "batch_in_epoch": batch_in_epoch, "batches_per_epoch": batches_per_epoch, "lr": float(self.optimizer.param_groups[0]["lr"]), "grad/norm": grad_norm, "amp/scale": float(self.scaler.get_scale()), "amp/dtype": self.config.training.amp_dtype, "optimizer/step": self.optimizer_step, "optimizer/updated": optimizer_updated, "optimizer/skipped": optimizer_skipped, "scene_ids": [scene_id for scene_id in batch.scene_ids], **_valid_ratios(batch.targets()), **scalar_values, **self.telemetry.snapshot(reset_window=True), **_memory_stats(self.device), **( { f"context_subset/{name}": value for name, value in context_subset_result.audit_diagnostics.items() } if context_subset_result is not None else {} ), **( {"context_subset": context_subset_metadata} if context_subset_metadata is not None else {} ), }, ) if ( step % self.config.training.checkpoint_every == 0 or step == self.config.training.max_steps ): latest_checkpoint = self.save_checkpoint() return latest_checkpoint def _learning_rate(self, optimizer_step: int) -> float: if self.config.training.lr_schedule == "constant": return self.config.training.learning_rate return _scheduled_learning_rate( self.config.training.learning_rate, optimizer_step, self.total_optimizer_steps, self.config.training.warmup_steps, self.config.training.min_lr_ratio, ) @torch.no_grad() def evaluate(self, checkpoint: str | Path | None = None) -> dict[str, float]: if checkpoint is not None: self.load_checkpoint(checkpoint, load_optimizer=False) self.model.eval() metric_records: list[dict[str, float]] = [] scene_records: dict[str, list[dict[str, float]]] = {} evaluation_loader = DataLoader( self.dataset, batch_size=self.config.training.batch_size, shuffle=False, num_workers=self.config.training.num_workers, collate_fn=collate_scene_examples, ) sample_count = 0 for batch in evaluation_loader: batch = batch.to(self.device) with torch.autocast( device_type=self.device.type, dtype=self.amp_dtype, enabled=self.amp_enabled, ): output = self.model( batch.context_rgb, batch.context_cameras, batch.target_cameras, batch.bounds ) targets = batch.targets() for batch_index in range(batch.context_rgb.shape[0]): predictions_item = { name: value[batch_index : batch_index + 1] for name, value in output.predictions.items() } targets_item = { name: value[batch_index : batch_index + 1] for name, value in targets.items() } metrics = compute_metrics( predictions_item, targets_item, bounds=batch.bounds[batch_index : batch_index + 1], ) scene_id = ( batch.scene_ids[batch_index] if batch.scene_ids and batch.scene_ids[batch_index] is not None else f"__sample_{sample_count:06d}" ) metric_records.append(metrics) scene_records.setdefault(scene_id, []).append(metrics) sample_count += 1 results = _legacy_window_macro(metric_records) per_scene_metrics = { scene_id: { "sample_count": len(records), "legacy_window_macro_metrics": _legacy_window_macro(records), "valid_count_weighted_metrics": _valid_count_weighted_metrics(records), "finite_window_counts": _finite_window_counts(records), } for scene_id, records in sorted(scene_records.items()) } report = { "schema_version": "mcss.evaluation.v1", "checkpoint": _checkpoint_metadata(self.loaded_checkpoint_path, self.step), "resolved_config": asdict(self.config), "evaluation": { "sample_count": sample_count, "scene_count": len(scene_records), }, "legacy_window_macro_metrics": results, "valid_count_weighted_metrics": _valid_count_weighted_metrics(metric_records), "per_scene_metrics": per_scene_metrics, "finite_window_counts": _finite_window_counts(metric_records), "aggregation_definitions": _aggregation_definitions(), } self.output_dir.mkdir(parents=True, exist_ok=True) _atomic_json(self.output_dir / "evaluation.json", results) _atomic_json(self.output_dir / "evaluation_report.json", report) return results def save_checkpoint(self) -> Path: self.checkpoint_dir.mkdir(parents=True, exist_ok=True) path = self._checkpoint_path(self.step) temporary = path.with_name(path.name + ".part") payload = { "step": self.step, "config": asdict(self.config), "model": self.model.state_dict(), "optimizer": self.optimizer.state_dict(), "scaler": self.scaler.state_dict(), "optimizer_step": self.optimizer_step, "telemetry": self.telemetry.state_dict(), "random_state": _random_state(), } if self.context_subset_generator is not None: payload["context_subset_generator_state"] = ( self.context_subset_generator.get_state() ) torch.save(payload, temporary) os.replace(temporary, path) self.loaded_checkpoint_path = path.resolve() return path def load_checkpoint(self, path: str | Path, *, load_optimizer: bool) -> None: checkpoint_path = Path(path).resolve() if not checkpoint_path.is_file(): raise FileNotFoundError(checkpoint_path) payload = torch.load(checkpoint_path, map_location=self.device, weights_only=False) self.model.load_state_dict(payload["model"]) self.step = int(payload.get("step", 0)) self.loaded_checkpoint_path = checkpoint_path if load_optimizer: self.optimizer.load_state_dict(payload["optimizer"]) self.optimizer_step = int( payload.get( "optimizer_step", self.step // self.config.training.gradient_accumulation ) ) if "scaler" in payload: self.scaler.load_state_dict(payload["scaler"]) if "telemetry" in payload: self.telemetry.load_state_dict(payload["telemetry"]) if "random_state" in payload: _restore_random_state(payload["random_state"]) if ( self.context_subset_generator is not None and "context_subset_generator_state" in payload ): self.context_subset_generator.set_state( payload["context_subset_generator_state"].cpu() ) def _checkpoint_path(self, step: int) -> Path: return self.checkpoint_dir / f"step_{step:06d}.pt" def build_dataset(config: DatasetConfig, *, seed: int) -> Dataset[SceneExample]: if config.name == "synthetic": return SyntheticSceneDataset( length=config.length, image_size=config.image_size, context_views=config.context_views, target_views=config.target_views, seed=seed, ) if config.root is None: raise ValueError("real datasets require a prepared manifest root") return ManifestSceneDataset( config.root, context_views=config.context_views, target_views=config.target_views, image_size=config.image_size, sample_stride=config.sample_stride, scene_ids=config.scene_ids, spatial_protocol=config.spatial_protocol, local_bounds_m=config.local_bounds_m, ) def resolve_device(requested: str) -> torch.device: if requested == "auto": return torch.device("cuda" if torch.cuda.is_available() else "cpu") device = torch.device(requested) if device.type == "cuda" and not torch.cuda.is_available(): raise RuntimeError("CUDA was requested but is not available") return device def seed_everything(seed: int) -> None: random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def _cycle(loader: DataLoader[SceneBatch]) -> Iterator[SceneBatch]: while True: yield from loader def _state_regularization( state: SceneState, *, evidence_residual_weight: float = 0.0, ) -> Tensor: density = torch.nn.functional.softplus(state.density_logits) sparsity = density.mean() * 1e-4 differences = [ (density[..., 1:] - density[..., :-1]).abs().mean(), (density[:, :, :, 1:, :] - density[:, :, :, :-1, :]).abs().mean(), (density[:, :, 1:, :, :] - density[:, :, :-1, :, :]).abs().mean(), ] regularization = sparsity + 1e-5 * sum(differences) if evidence_residual_weight == 0.0: return regularization if state.transport_evidence is not None: transport = state.transport_evidence density_rewrite = ( transport.surface_localization_support * transport.density_residual.abs() ) if state.appearance is None: color_rewrite = ( transport.appearance_confidence * transport.color_logit_residual.abs().mean(dim=1, keepdim=True) ) else: appearance = state.appearance color_rewrite = ( appearance.confidence * appearance.color_logit_residual.abs().mean(dim=1, keepdim=True) ) elif state.dual_evidence is not None: dual = state.dual_evidence density_rewrite = ( dual.surface_localization_support * dual.density_residual.abs() ) if state.appearance is None: color_rewrite = ( dual.appearance_confidence * dual.color_logit_residual.abs().mean(dim=1, keepdim=True) ) else: appearance = state.appearance color_rewrite = ( appearance.confidence * appearance.color_logit_residual.abs().mean(dim=1, keepdim=True) ) elif state.evidence is not None: evidence = state.evidence density_rewrite = evidence.confidence * evidence.density_residual.abs() if state.appearance is None: color_rewrite = evidence.confidence * evidence.color_logit_residual.abs().mean( dim=1, keepdim=True ) else: appearance = state.appearance color_rewrite = appearance.confidence * appearance.color_logit_residual.abs().mean( dim=1, keepdim=True ) else: return regularization rewrite_penalty = density_rewrite.mean() + color_rewrite.mean() return regularization + evidence_residual_weight * rewrite_penalty def _state_diagnostics(state: SceneState) -> dict[str, float]: if state.transport_evidence is not None: transport = state.transport_evidence diagnostics = { "state/surface_localization_support_mean": float( transport.surface_localization_support.detach().mean().item() ), "state/transported_appearance_confidence_mean": float( transport.transported_appearance_confidence.detach().mean().item() ), "state/appearance_confidence_mean": float( transport.appearance_confidence.detach().mean().item() ), "state/surface_peakness_mean": float( transport.surface_peakness.detach().mean().item() ), "state/transport_coverage_mean": float( transport.transport_coverage.detach().mean().item() ), "state/transport_coverage_nonzero_fraction": float( (transport.transport_coverage.detach() > 0).float().mean().item() ), "state/localization_unknown_mean": float( transport.localization_unknown_probability.detach().mean().item() ), "state/appearance_unknown_mean": float( transport.appearance_unknown_probability.detach().mean().item() ), "state/localization_completion_gate_mean": float( transport.localization_completion_gate.detach().mean().item() ), "state/appearance_completion_gate_mean": float( transport.appearance_completion_gate.detach().mean().item() ), "state/density_residual_abs": float( transport.density_residual.detach().abs().mean().item() ), "state/color_residual_abs": float( transport.color_logit_residual.detach().abs().mean().item() ), "state/surface_peakness_nonzero_fraction": float( (transport.surface_peakness.detach() > 0).float().mean().item() ), "state/localization_density_rewrite_mean": float( ( transport.surface_localization_support.detach() * transport.density_residual.detach().abs() ).mean().item() ), "state/appearance_color_rewrite_mean": float( ( transport.appearance_confidence.detach() * transport.color_logit_residual.detach().abs().mean( dim=1, keepdim=True ) ).mean().item() ), } diagnostics.update( _decile_diagnostics( "state/surface_localization_support", transport.surface_localization_support, ) ) diagnostics.update( _decile_diagnostics( "state/transported_appearance_confidence", transport.transported_appearance_confidence, ) ) diagnostics.update( _decile_diagnostics("state/surface_peakness", transport.surface_peakness) ) diagnostics.update( _decile_diagnostics("state/transport_coverage", transport.transport_coverage) ) elif state.dual_evidence is not None: dual = state.dual_evidence diagnostics = { "state/surface_localization_support_mean": float( dual.surface_localization_support.detach().mean().item() ), "state/appearance_confidence_mean": float( dual.appearance_confidence.detach().mean().item() ), "state/surface_peakness_mean": float( dual.surface_peakness.detach().mean().item() ), "state/localization_unknown_mean": float( dual.localization_unknown_probability.detach().mean().item() ), "state/appearance_unknown_mean": float( dual.appearance_unknown_probability.detach().mean().item() ), "state/localization_completion_gate_mean": float( dual.localization_completion_gate.detach().mean().item() ), "state/appearance_completion_gate_mean": float( dual.appearance_completion_gate.detach().mean().item() ), "state/density_residual_abs": float( dual.density_residual.detach().abs().mean().item() ), "state/color_residual_abs": float( dual.color_logit_residual.detach().abs().mean().item() ), "state/surface_peakness_nonzero_fraction": float( (dual.surface_peakness.detach() > 0).float().mean().item() ), "state/localization_density_rewrite_mean": float( ( dual.surface_localization_support.detach() * dual.density_residual.detach().abs() ).mean().item() ), "state/appearance_color_rewrite_mean": float( ( dual.appearance_confidence.detach() * dual.color_logit_residual.detach().abs().mean(dim=1, keepdim=True) ).mean().item() ), } diagnostics.update( _decile_diagnostics( "state/surface_localization_support", dual.surface_localization_support, ) ) diagnostics.update( _decile_diagnostics( "state/appearance_confidence", dual.appearance_confidence, ) ) diagnostics.update( _decile_diagnostics("state/surface_peakness", dual.surface_peakness) ) elif state.evidence is not None: evidence = state.evidence diagnostics = { "state/evidence_mean": float(evidence.confidence.detach().mean().item()), "state/unknown_mean": float(evidence.unknown_probability.detach().mean().item()), "state/completion_gate_mean": float(evidence.completion_gate.detach().mean().item()), "state/density_residual_abs": float( evidence.density_residual.detach().abs().mean().item() ), "state/color_residual_abs": float( evidence.color_logit_residual.detach().abs().mean().item() ), } else: return {} if state.appearance is not None: appearance = state.appearance if state.dual_evidence is None and state.transport_evidence is None: diagnostics.update({ "state/appearance_evidence_mean": float( appearance.confidence.detach().mean().item() ), "state/appearance_unknown_mean": float( appearance.unknown_probability.detach().mean().item() ), "state/appearance_completion_gate_mean": float( appearance.completion_gate.detach().mean().item() ), "state/appearance_color_residual_abs": float( appearance.color_logit_residual.detach().abs().mean().item() ), "state/appearance_resolution_scale": ( appearance.spatial_shape[0] / state.spatial_shape[0] ), }) else: diagnostics.update({ "state/highres_appearance_confidence_mean": float( appearance.confidence.detach().mean().item() ), "state/highres_appearance_unknown_mean": float( appearance.unknown_probability.detach().mean().item() ), "state/highres_appearance_completion_gate_mean": float( appearance.completion_gate.detach().mean().item() ), "state/highres_appearance_color_residual_abs": float( appearance.color_logit_residual.detach().abs().mean().item() ), "state/appearance_resolution_scale": ( appearance.spatial_shape[0] / state.spatial_shape[0] ), }) return diagnostics def _decile_diagnostics(prefix: str, values: Tensor) -> dict[str, float]: probabilities = torch.arange( 0.1, 1.0, 0.1, device=values.device, dtype=torch.float32, ) quantiles = torch.quantile(values.detach().float(), probabilities) return { f"{prefix}_p{index * 10:02d}": float(value.item()) for index, value in enumerate(quantiles, start=1) } def _legacy_window_macro(records: list[dict[str, float]]) -> dict[str, float]: names = sorted({name for record in records for name in record}) results: dict[str, float] = {} for name in names: values = [record[name] for record in records if name in record] if any(math.isinf(value) for value in values): results[name] = float("inf") continue finite = [value for value in values if math.isfinite(value)] if finite: results[name] = sum(finite) / len(finite) return results def _finite_window_counts(records: list[dict[str, float]]) -> dict[str, int]: names = sorted({name for record in records for name in record}) return { name: sum(math.isfinite(record[name]) for record in records if name in record) for name in names } def _valid_count_weighted_metrics(records: list[dict[str, float]]) -> dict[str, float]: names = sorted({name for record in records for name in record}) count_names = { name for name in names if name.endswith("/valid_pixels") or name.endswith("/valid_points") or name.endswith("/total_pixels") } results = { name: sum( record[name] for record in records if name in record and math.isfinite(record[name]) ) for name in sorted(count_names) } for name in names: if name in count_names or name in {"rgb/psnr", "support/coverage"}: continue weight_name = _metric_weight_name(name) if weight_name is None: continue weighted_values = [ (record[name], record[weight_name]) for record in records if name in record and weight_name in record and math.isfinite(record[name]) and math.isfinite(record[weight_name]) and record[weight_name] > 0 ] total_weight = sum(weight for _, weight in weighted_values) if total_weight == 0: continue if name in {"depth/rmse", "point/rmse_m"}: value = math.sqrt( sum(metric**2 * weight for metric, weight in weighted_values) / total_weight ) else: value = sum(metric * weight for metric, weight in weighted_values) / total_weight results[name] = value if "rgb/mse" in results: mse = results["rgb/mse"] results["rgb/psnr"] = float("inf") if mse == 0.0 else -10.0 * math.log10(mse) valid_support = results.get("support/valid_pixels") total_support = results.get("support/total_pixels") if valid_support is not None and total_support: results["support/coverage"] = valid_support / total_support return dict(sorted(results.items())) def _metric_weight_name(name: str) -> str | None: prefix = name.split("/", maxsplit=1)[0] return { "rgb": "rgb/valid_pixels", "depth": "depth/valid_pixels", "normal": "normal/valid_pixels", "point": "point/valid_points", "visibility": "visibility/valid_pixels", }.get(prefix) def _checkpoint_metadata(path: Path | None, step: int) -> dict[str, object]: if path is None: return {"absolute_path": None, "sha256": None, "step": step} return { "absolute_path": str(path.resolve()), "sha256": _sha256_file(path), "step": step, } def _sha256_file(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as handle: while chunk := handle.read(1024 * 1024): digest.update(chunk) return digest.hexdigest() def _aggregation_definitions() -> dict[str, object]: return { "legacy_window_macro": ( "Arithmetic mean over finite per-sample window metrics; any infinite perfect-score " "window keeps that metric infinite for evaluation.json compatibility." ), "valid_count_weighted": ( "Per-window metrics weighted by their valid pixel or point counts; count fields are " "summed across windows." ), "pooled_derivations": { "rgb/psnr": "Derived from the valid-count-pooled RGB MSE.", "depth/rmse": "Square root of pooled valid depth squared error.", "point/rmse_m": "Square root of pooled valid point squared error.", "support/coverage": "Total supported query pixels divided by total query pixels.", }, "non_additive_metrics": ( "SSIM, normal thresholds, Chamfer, F-scores, visibility IoU, and visibility F1 are " "computed within each window and then valid-count weighted. Point-cloud nearest " "neighbors are never matched across different scenes." ), } def _append_jsonl(path: Path, record: dict[str, object]) -> None: path.parent.mkdir(parents=True, exist_ok=True) with path.open("a", encoding="utf-8") as handle: handle.write(json.dumps(record, sort_keys=True) + "\n") def _atomic_json(path: Path, data: object) -> None: path.parent.mkdir(parents=True, exist_ok=True) temporary = path.with_name(path.name + ".part") temporary.write_text(json.dumps(data, indent=2, sort_keys=True) + "\n", encoding="utf-8") os.replace(temporary, path) def _random_state() -> dict[str, object]: return { "python": random.getstate(), "numpy": np.random.get_state(), "torch": torch.get_rng_state(), "cuda": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None, } def _restore_random_state(state: dict[str, object]) -> None: random.setstate(state["python"]) # type: ignore[arg-type] np.random.set_state(state["numpy"]) # type: ignore[arg-type] torch_state = state["torch"] if not isinstance(torch_state, Tensor): raise TypeError("torch RNG state must be a tensor") torch.set_rng_state(torch_state.detach().cpu().to(dtype=torch.uint8)) if state.get("cuda") is not None and torch.cuda.is_available(): cuda_states = state["cuda"] if not isinstance(cuda_states, (list, tuple)) or not all( isinstance(cuda_state, Tensor) for cuda_state in cuda_states ): raise TypeError("CUDA RNG state must be a sequence of tensors") torch.cuda.set_rng_state_all( [cuda_state.detach().cpu().to(dtype=torch.uint8) for cuda_state in cuda_states] ) def _scheduled_learning_rate( base_learning_rate: float, optimizer_step: int, total_optimizer_steps: int, warmup_steps: int, min_lr_ratio: float, ) -> float: if base_learning_rate <= 0: raise ValueError("base_learning_rate must be positive") if total_optimizer_steps < 1 or not 1 <= optimizer_step <= total_optimizer_steps: raise ValueError("optimizer_step must be within the configured optimizer updates") if not 0 <= warmup_steps <= total_optimizer_steps: raise ValueError("warmup_steps must fit within total_optimizer_steps") if not 0.0 <= min_lr_ratio <= 1.0: raise ValueError("min_lr_ratio must be between zero and one") if warmup_steps and optimizer_step <= warmup_steps: return base_learning_rate * optimizer_step / warmup_steps decay_steps = total_optimizer_steps - warmup_steps if decay_steps == 0: return base_learning_rate progress = (optimizer_step - warmup_steps) / decay_steps cosine = 0.5 * (1.0 + math.cos(math.pi * progress)) return base_learning_rate * (min_lr_ratio + (1.0 - min_lr_ratio) * cosine) def _float_mapping(value: object, name: str) -> dict[str, float]: if not isinstance(value, Mapping): raise TypeError(f"telemetry {name} must be a mapping") return {str(key): float(item) for key, item in value.items()} def _valid_ratios(targets: Mapping[str, Tensor]) -> dict[str, float]: support = targets.get("support") support_valid = None if support is None else torch.isfinite(support) & (support > 0.5) ratios = { "valid/support": 1.0 if support_valid is None else float(support_valid.float().mean()) } for name, target in targets.items(): if name == "support": continue if name in {"rgb", "normal", "point"}: valid = torch.isfinite(target).all(dim=-3, keepdim=True) else: valid = torch.isfinite(target) if name == "depth": valid = valid & (target > 0) elif name == "normal": valid = valid & (target.square().sum(dim=-3, keepdim=True) > 1e-8) elif name == "point" and "visibility" in targets: valid = valid & (targets["visibility"] > 0.5) if support_valid is not None: valid = valid & support_valid ratios[f"valid/{name}"] = float(valid.float().mean().item()) return ratios def _memory_stats(device: torch.device) -> dict[str, float]: if device.type != "cuda": return {"memory/max_allocated_mb": 0.0, "memory/max_reserved_mb": 0.0} scale = 1024.0**2 return { "memory/max_allocated_mb": torch.cuda.max_memory_allocated(device) / scale, "memory/max_reserved_mb": torch.cuda.max_memory_reserved(device) / scale, }