gaussian-optimizer-2d

A minimal PyTorch project for optimizing a 2D Gaussian scene against a procedurally generated ground truth. Built as a learning tool to understand differentiable rendering and gradient-based scene optimization before stepping into full 3D Gaussian Splatting.

Features

• 2D Gaussian scene model — learnable parameters per Gaussian: position, scale, rotation, opacity, and color.
• Differentiable renderer — rasterizes Gaussians onto a pixel grid using weighted-average alpha blending; fully vectorised, no sequential loop.
• Ground truth generator — renders a fixed random scene and exposes it as both a training tensor and a saveable PIL image.
• Initializer with difficulty levels — initialize the scene at four levels relative to the ground truth, from trivial (EXACT) to hard (RANDOM), to stress-test optimizer convergence from different starting points.
• L1 + SSIM loss — combined perceptual loss following the 3DGS paper convention (α = 0.15 L1, 0.85 SSIM).
• Learning rate scheduler — cosine annealing, exponential decay, or constant; configurable via LRSchedule enum.
• Adaptive densification — clones under-sampled Gaussians, splits over-sized ones, and prunes near-transparent ones; Adam momentum is preserved for surviving Gaussians across densification steps.
• Training timing — per-step wall times logged at every interval, plus a full summary (total, mean, min, max ms/step, throughput) attached to every TrainResult for easy performance comparison.
• Device auto-detection — selects CUDA → MPS → CPU automatically; torch.compile is applied on CUDA for additional throughput.
• Notebooks — interactive demos for ground truth generation and training visualization across init levels.

Project structure

gaussian_optimizer_2d/
├── src/
│   └── gaussian_optimizer_2d/
│       ├── gaussian2d.py       # Learnable Gaussian2DScene (nn.Module)
│       ├── ground_truth.py     # GroundTruth dataclass + generator
│       ├── initializer.py      # InitLevel enum + initialize_scene()
│       ├── renderer.py         # render() — weighted-average rasterizer
│       ├── loss.py             # L1SSIMLoss
│       ├── densification.py    # DensifyConfig, DensificationState, densify()
│       └── train.py            # TrainConfig, TrainResult, TimingStats, train()
├── notebooks/
│   ├── ground_truth_demo.ipynb   # Generate and display a ground truth scene
│   └── training_demo.ipynb       # Train across init levels, compare results
├── tests/
├── docs/
│   └── rendering_blending_modes.md  # Analysis of blending approaches
├── pyproject.toml
└── Makefile

Setup

make install

For notebooks, install dev dependencies:

uv sync --extra dev

Usage

Run training from the command line:

make run

Use the API:

from src.gaussian_optimizer_2d.train import train, TrainConfig, LRSchedule
from src.gaussian_optimizer_2d.initializer import InitLevel
from src.gaussian_optimizer_2d.densification import DensifyConfig

result = train(TrainConfig(
    n_gaussians=20,
    height=256,
    width=256,
    n_iterations=1000,
    lr=1e-2,
    lr_schedule=LRSchedule.COSINE,
    init_level=InitLevel.RANDOM,
    densify=DensifyConfig(every_n_steps=100, max_gaussians=500),
))

print(result.timing.summary())
result.ground_truth.save("ground_truth.png")

Open a notebook:

uv run jupyter notebook notebooks/training_demo.ipynb

Init levels

Level	Description
EXACT	Copy ground truth params — loss ≈ 0 from step 0
NOISY	Perturb true params with small noise
POSITIONS_ONLY	True Gaussian centers, random appearance
RANDOM	Fully random, independent of ground truth

Densification

Adaptive density control runs every every_n_steps steps:

Condition	Action
High gradient + small scale	Clone — add a nearby copy to increase coverage
High gradient + large scale	Split — replace with two smaller children along the principal axis
Opacity below threshold	Prune — remove the Gaussian unconditionally

Adam momentum is preserved for surviving and cloned Gaussians; split children start fresh.

Training timing

Every TrainResult includes a TimingStats object:

result.timing.total_s        # total wall time in seconds
result.timing.mean_ms        # average ms per step
result.timing.steps_per_sec  # throughput
result.timing.summary()      # one-line string for quick comparison

Per-interval timing is also printed during training alongside loss and learning rate.