This repository builds on the official implementation of the paper What's in a Prior? Learned Proximal Networks for Inverse Problems @ ICLR 2024 by Zhenghan Fang, Sam Buchanan, and Jeremias Sulam.
We adapt their Learned Proximal Network (LPN) framework to investigate distribution-shift uncertainty estimation in inverse problems.
Our experiments focus specifically on MNIST-based sparse-view CT reconstruction.
👉 This repository is an extension of their work; please see the original repo for the full method and additional experiments.
👉 All pre-trained models and results are provided in this repository, so the evaluate_ct.ipynb notebook file can be executed directly without retraining.
The code is implemented with Python 3.9.16 and PyTorch 1.12.0.
Install the conda environment:
conda env create -f environment.ymlInstall the lpn package:
pip install -e .The datasets are placed in the data/ folder.
The dataset is located in data/mayoct/mnist.
The only0 subset is used for in-distribution training, while idood_valid is used for evaluating distribution-shift performance (in- vs out-of-distribution).
Code of the main functionalities of LPN is placed in the lpn folder.
Code for reproducing the MNIST experiments is placed in the exps/mayoct folder.
- Train and Compute prior:
bash run.sh # trains the Learned Proximal Network (LPN)
bash run_ct.sh # performs CT reconstruction using the trained prior- Trained model:
exps/mayoct/models_mnist/lpn - Training logs:
exps/mayoct/models_mnist - CT results:
exps/mayoct/results/inverse/mayoct/tomo
- Visualize results
- Evaluation and visualization available at
evaluate_ct.ipynb. All pre-trained models and results are provided in this repository, so the notebook can be executed directly without retraining.
This work builds directly on:
- What's in a Prior? Learned Proximal Networks for Inverse Problems (ICLR 2024)
- Original implementation
Additional references: