{"id":25911743,"url":"https://github.com/devzhk/enkg-pytorch","last_synced_at":"2026-05-12T14:31:54.263Z","repository":{"id":279703888,"uuid":"939575571","full_name":"devzhk/enkg-pytorch","owner":"devzhk","description":"Offical Pytorch Implementation of EnKG","archived":false,"fork":false,"pushed_at":"2025-02-26T23:32:19.000Z","size":68,"stargazers_count":0,"open_issues_count":0,"forks_count":0,"subscribers_count":1,"default_branch":"main","last_synced_at":"2025-02-27T00:30:14.188Z","etag":null,"topics":["derivative-free","diffusion-model","inverse-problem"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/devzhk.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2025-02-26T19:01:10.000Z","updated_at":"2025-02-26T23:32:22.000Z","dependencies_parsed_at":"2025-02-27T00:30:18.649Z","dependency_job_id":"5d86544b-2cfd-41b6-bb16-be84c682a2f9","html_url":"https://github.com/devzhk/enkg-pytorch","commit_stats":null,"previous_names":["devzhk/enkg-pytorch"],"tags_count":3,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/devzhk%2Fenkg-pytorch","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/devzhk%2Fenkg-pytorch/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/devzhk%2Fenkg-pytorch/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/devzhk%2Fenkg-pytorch/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/devzhk","download_url":"https://codeload.github.com/devzhk/enkg-pytorch/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":241637264,"owners_count":19994946,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["derivative-free","diffusion-model","inverse-problem"],"created_at":"2025-03-03T09:19:05.090Z","updated_at":"2026-05-12T14:31:54.227Z","avatar_url":"https://github.com/devzhk.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"## Ensemble Kalman Diffusion Guidance - Offical Pytorch Implementation\n[Ensemble Kalman Diffusion Guidance: A Derivative-free Method for Inverse Problems](https://arxiv.org/abs/2409.20175)\n\nHongkai Zheng, Wenda Chu*, Austin Wang*, Nikola Kovachki, Ricardo Baptista, Yisong Yue\n\nAbstract: \n*When solving inverse problems, one increasingly popular approach is to use pre-trained diffusion models as plug-and-play priors. This framework can accommodate different forward models without re-training while preserving the generative capability of diffusion models.  Despite their success in many imaging inverse problems, most existing methods rely on privileged information such as derivative, pseudo-inverse, or full knowledge about the forward model. This reliance poses a substantial limitation that restricts their use in a wide range of problems where such information is unavailable, such as in many scientific applications. We propose Ensemble Kalman Diffusion Guidance (EnKG), a derivative-free approach that can solve inverse problems by only accessing forward model evaluations and a pre-trained diffusion model prior. We study the empirical effectiveness of EnKG across various inverse problems, including scientific settings such as inferring fluid flows and astronomical objects, which are highly non-linear inverse problems that often only permit black-box access to the forward model.*\n\n\n## Environment requirements\n- We recommend Linux with 64-bit Python 3.11.5 and Pytorch 2.2.2. See https://pytorch.org for PyTorch install instructions.\n- `torch, accelerate, hydra-core, ehtim, ehtplot, piq, wandb, pillow, lmbd, omegaconf` are the main Python libraries required. Environment file is provided in `env.yml`. \n- We also provide a [Dockerfile](Docker/Dockerfile) under `Docker`. You can use as follows:\n\n```bash\n# Build docker image\ndocker build -t [image tag] --build-arg USER_ID=$(id -u) --build-arg GROUP_ID=$(id -g) .\n\n# Run docker container\ndocker run --gpus all -it --rm --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 -v [path to the top of this git repo]:/enkg -v [path to data]:/data [image tag]\n```\nBreakdown of the `docker run` command:\n- `--gpus all -it --rm`: Enable all available GPUs, starts an interactive session, and automatically remove the container upon exit.\n- `--ipc=host --ulimit memlock=-1 --ulimit stack=67108864`: Recommended NVIDIA flags to unlock resource constraints.\n- `-v [path to the top of this repo]:/enkg -v [path to data]:/data`: Mount the current dir to `/enkg`. Mount the data directory to `/data`.\n\n## Data and pre-trained models\nData and pre-trained models can be found in the Github release page. By default, the data should be placed in `../data` and the pre-trained models should be placed in `checkpoints` directory. You can also specify the data and checkpoint path in the config file.\n\n## Inference\n\nBy default, `configs/config.yaml` will be loaded for inference. You can override the config value by\n```bash\npython3 main.py problem=[inverse problem config name] algorithm=[algorithm config name] pretrain=[pretrained model config name]\n```\nThe structure of the inference config is explained below. \n| Key       | Description                                                                      |\n|-----------|----------------------------------------------------------------------------------|\n| `problem`   | Name of the inverse problem configuration. (See `configs/problem`)             |\n| `algorithm` | Name of the algorithm configuration. (See `configs/algorithm`)                 |\n| `pretrain`  | Name of the pre-trained model configuration. (see `configs/pretrain`)          |\n| `tf32`      | (bool) Enables TF32 mode for improved performance on Ampere+ GPUs.             |\n| `compile`   | (bool) Enable `torch.compile` (recommended for ensemble methods).              |\n| `seed`      | (int) Random seed.                                                             |\n| `inference` | (bool) If False, skip inference and only run evaluation.                       |\n| `exp_name`  | (string) Sets the experiment name for logging and saving results.              |\n| `wandb`     | (bool) Enables logging to Weights \u0026 Biases (WandB).                            |\n\nWe provide sample scripts to run experiments in `scripts`. \n- `scripts/navier-stokes.sh` contains commands to run different algorithms on the inverse problem of the Navier-Stokes equation. (Takes ~2 hours on an A100 GPU as the numerical solver takes time to run) \n- `scripts/ffhq.sh` contains commands to run different algorithms on image restoration tasks. In general, image restoration tasks here are not the best use case for derivative-free methods. For example, EnKG is inefficient when the forward model is much faster than diffusion model evaluation. This serves as a proof-of-concept example. \n- `scripts/blackhole.sh` contains commands to run inference on black hole imaging tasks. (These experiments run efficiently on an A100 GPU.)\n\n\n## License\nThis project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.\n\n\n## Citation\n```bibtex   \n@article{zheng2024ensemble,\n  title={Ensemble kalman diffusion guidance: A derivative-free method for inverse problems},\n  author={Zheng, Hongkai and Chu, Wenda and Wang, Austin and Kovachki, Nikola and Baptista, Ricardo and Yue, Yisong},\n  journal={arXiv preprint arXiv:2409.20175},\n  year={2024}\n}\n```\n\n## Acknowledgements\n- The pre-trained model weights for FFHQ256 is converted from [DPS's repository](https://github.com/DPS2022/diffusion-posterior-sampling). We thank the authors for releasing their pre-trained model. \n- We thank Ben Prather, Abhishek Joshi, Vedant Dhruv, C.K. Chan, and Charles Gammie for\nthe synthetic blackhole images [GRMHD Dataset](https://iopscience.iop.org/article/10.3847/1538-4365/ac582e) used here, generated under NSF grant AST 20-34306. ","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fdevzhk%2Fenkg-pytorch","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fdevzhk%2Fenkg-pytorch","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fdevzhk%2Fenkg-pytorch/lists"}