How akhasanovv/admm-implementation works
Purpose and entry points
This is a PyTorch reconstruction system for mask-based lensless images, implementing fixed ADMM, trainable unrolled Le-ADMM, and modular pre/post-processing variants. Primary executables are train.py, inference.py, calculate_metrics.py, and benchmark_speed.py; demo.ipynb is a Colab-oriented demonstration.
Configuration-driven composition
Runtime behavior is assembled from hierarchical YAML configs: model, dataset split, dataloader, transforms, metrics, and experiment writer are independently selectable/overridable from the CLI. Core defaults and inference composition live in src/configs/baseline.yaml, src/configs/inference.yaml, and model presets under src/configs/model/.
Data ingestion and physical inputs
Datasets provide lensless measurements, per-image mask/PSF arrays, and optional lensed ground truth. The DigiCam dataset and custom directory layout are handled through src/datasets/lensless.py, src/datasets/base_dataset.py, and dataset profiles in src/configs/datasets/. Batch construction is separated in src/datasets/collate.py; lensless-specific preprocessing and optical/sensor representations reside in src/lensless_helpers/preprocessor.py, src/lensless_helpers/psf.py, and src/lensless_helpers/sensor.py.
Reconstruction model family
The reconstruction boundary is the model package. Classical fixed-parameter ADMM is in src/model/admm.py; learned unrolled ADMM is in src/model/le_admm.py; modular learned pre-/post-processing composition is in src/model/modular_le_admm.py. The learned denoising/prior network used by these methods is src/model/drunet.py, while src/model/baseline_model.py supplies the base model contract.
Training and checkpoint lifecycle
train.py constructs the configured pipeline and delegates optimization lifecycle to src/trainer/trainer.py, extending shared control logic in src/trainer/base_trainer.py. Training consumes reconstructed outputs and targets through configurable losses, notably src/loss/admm_loss.py, and writes model checkpoints under the configured run directory. Fixed ADMM remains usable without a training stage.
Inference, persistence, and metrics
inference.py uses src/trainer/inferencer.py to load optional pretrained weights, execute reconstruction over a split, and save predictions. Metrics are pluggable: reconstruction quality calculations and aggregation are centered in src/metrics/reconstruction.py, src/metrics/base_metric.py, and src/metrics/tracker.py. calculate_metrics.py supports post hoc comparison of saved predictions against external ground truth directories.
External services and runtime dependencies
The material runtime is Python with PyTorch and its configured package dependencies in requirements.txt. Experiment telemetry can target Comet ML or Weights & Biases via src/logger/cometml.py, src/logger/wandb.py, and src/logger/logger.py; Comet requires an API key. Pretrained model artifacts are externally hosted on Hugging Face, while the demo also retrieves data/checkpoints from external sources.
Main entry points: trainer.py Trainer · baseline.yaml YAML composition · inferencer.py Inferencer · le_admm.py Unrolled Le-ADMM
Components 21
Entrypoints
- train.py
train.py - inference.py
inference.py - calculate_metrics.py · Compares saved predictions with external ground truth.
calculate_metrics.py
Configuration
- YAML composition · CLI-overridable model, data, transforms, metrics, and writer selection.
src/configs/baseline.yaml - Inference config
src/configs/inference.yaml
Data & optics
- Lensless datasets · Loads measurements, per-image mask/PSF, and optional lensed targets.
src/datasets/lensless.py - Batch collation
src/datasets/collate.py - Optical preprocessing · Prepares lensless inputs and PSF/sensor representations.
src/lensless_helpers/preprocessor.py
Reconstruction models
- Base model contract
src/model/baseline_model.py - Fixed ADMM · Runs fixed-parameter reconstruction without training.
src/model/admm.py - Unrolled Le-ADMM
src/model/le_admm.py - Modular Le-ADMM · Adds learned pre- and post-processing around unrolled ADMM.
src/model/modular_le_admm.py - DRUNet prior
src/model/drunet.py
Training & inference
- Trainer · Optimizes configured trainable reconstruction models and writes checkpoints.
src/trainer/trainer.py - Inferencer · Loads optional weights, reconstructs a split, and saves predictions.
src/trainer/inferencer.py - ADMM loss
src/loss/admm_loss.py
Evaluation & services
- Reconstruction metrics · Computes and aggregates reconstruction quality.
src/metrics/reconstruction.py - Saved predictions
- Run checkpoints
- Comet ML / W&B · Optional external experiment logging.
src/logger/logger.py - Hugging Face · External host for pretrained model artifacts.
Connections 29
- train.py → YAML composition: composes
- inference.py → Inference config: composes
- YAML composition → Lensless datasets: selects split
- YAML composition → Fixed ADMM: selects model
- YAML composition → Unrolled Le-ADMM: selects model
- YAML composition → Modular Le-ADMM: selects model
- Lensless datasets → Batch collation: samples
- Batch collation → Optical preprocessing: batches inputs
- Optical preprocessing → Fixed ADMM: measurements and PSF
- Optical preprocessing → Unrolled Le-ADMM: measurements and PSF
- Optical preprocessing → Modular Le-ADMM: measurements and PSF
- Base model contract → Fixed ADMM: implements
- Base model contract → Unrolled Le-ADMM: implements
- Base model contract → Modular Le-ADMM: implements
- DRUNet prior → Unrolled Le-ADMM: learned prior
- DRUNet prior → Modular Le-ADMM: learned prior
- train.py → Trainer: delegates
- Trainer → Unrolled Le-ADMM: optimizes
- Trainer → Modular Le-ADMM: optimizes
- ADMM loss → Trainer: supervises
- Trainer → Run checkpoints: writes
- Trainer → Comet ML / W&B: logs to
- inference.py → Inferencer: delegates
- Run checkpoints → Inferencer: optional weights
- Hugging Face → Inferencer: pretrained weights
- Inferencer → Saved predictions: saves
- Saved predictions → calculate_metrics.py: compares
- calculate_metrics.py → Reconstruction metrics: calculates
- Inferencer → Reconstruction metrics: evaluates