{"id":21730607,"url":"https://github.com/dlr-mi/atme","last_synced_at":"2025-10-10T09:12:06.804Z","repository":{"id":163030841,"uuid":"624373475","full_name":"DLR-MI/atme","owner":"DLR-MI","description":"Official pytorch implementation of ATME [CVPR 2023 GMCV]","archived":false,"fork":false,"pushed_at":"2023-11-02T10:31:37.000Z","size":34750,"stargazers_count":29,"open_issues_count":2,"forks_count":2,"subscribers_count":2,"default_branch":"main","last_synced_at":"2025-06-08T00:45:51.075Z","etag":null,"topics":["deep-learning","diffusion-models","gan","image-translation","pix2pix"],"latest_commit_sha":null,"homepage":"","language":"Jupyter Notebook","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"other","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/DLR-MI.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}},"created_at":"2023-04-06T10:22:20.000Z","updated_at":"2025-02-26T06:38:26.000Z","dependencies_parsed_at":null,"dependency_job_id":"81552058-f646-4114-bc4e-a23bd6a8ad6e","html_url":"https://github.com/DLR-MI/atme","commit_stats":null,"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/DLR-MI/atme","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/DLR-MI%2Fatme","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/DLR-MI%2Fatme/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/DLR-MI%2Fatme/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/DLR-MI%2Fatme/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/DLR-MI","download_url":"https://codeload.github.com/DLR-MI/atme/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/DLR-MI%2Fatme/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":279003389,"owners_count":26083579,"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","status":"online","status_checked_at":"2025-10-10T02:00:06.843Z","response_time":62,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"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":["deep-learning","diffusion-models","gan","image-translation","pix2pix"],"created_at":"2024-11-26T04:16:24.724Z","updated_at":"2025-10-10T09:12:06.787Z","avatar_url":"https://github.com/DLR-MI.png","language":"Jupyter Notebook","funding_links":[],"categories":[],"sub_categories":[],"readme":"\n\u003c!-- # ATME --\u003e\n# \u003cimg src=\"imgs/atme.png\" alt=\"ATME\" width=\"100\" height=\"50\"\u003e \n[ArXiv](https://arxiv.org/abs/2304.09024) |\n[Medium](https://link.medium.com/pOw7BBY2Ezb) \n\n\nThis repository is the official implementation of [Look ATME: The Discriminator Mean Entropy Needs Attention](https://arxiv.org/abs/2304.09024), accepted for the CVPR 2023 [Workshop on Generative Models for Computer Vision](https://generative-vision.github.io/workshop-CVPR-23/).\n\n\u003cimg src='imgs/mosaic_noisy.png'\u003e \n\n\u003cbr\u003e\n\nATME is a GAN where, for each iteration at epoch $t$, the input images for the generator are corrupted with a noisy representation of the discriminator’s decision map at epoch $t − 1$. The adversarial game enforces this noise to be removed, sending the proper signals for equilibration by encouraging the discriminator to converge towards its maximum entropy state.\n\nThe implementation is based on the official [pix2pix](https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix) pytorch code. Since it uses the UNet for denoising diffusion probabilistic models as generator, the corresponding module is borrowed from [lucidrains](https://github.com/lucidrains/denoising-diffusion-pytorch/blob/main/README.md).\n\n## Requirements\nFirst clone this repository:\n```\ngit clone https://github.com/DLR-MI/atme.git\ncd atme\n```\nand install suitable dependencies:\n- For pip users, you need to `pip install -r requirements.txt`.\n- For conda users `conda env create -f environment.yml` and `conda activate atme`.\n\n\n## Datasets\nTo run the experiments in the paper, you first need to download the respective datasets. This can be done by running \n```\nbash ./datasets/download_pix2pix_dataset.sh DATASET\n```\nwhere `DATASET` is any of `facades`, `maps`, `cityscapes`, or `night2day`. The result is placed in the dataroot folder `./datasets/DATASET`.\n\n## Training\nAfter downloading the dataset, decide on a project name (say `PROJECT`) and run\n\n```\npython train.py --name PROJECT --model atme --batch_size 48 --direction AtoB --dataroot ./datasets/DATASET\n```\nFor multi-GPU training, add the option `--gpu_ids 0,1,2,3,...` to the command above.\n\nIn order to visualize intermediate results during the training process, you can either:\n- run `python -m visdom.server` in a separate terminal, before the command above, and visit http://localhost:8097 afterwards, or\n- check out `./checkpoints/PROJECT/web/index.html`. \n\nYou should then see snapshots corresponding to $x$, $W_{t}$, $x_t$, $\\hat{y}$, and $y$ in the paper, and named with the suffices real_A, Disc_B, noisy_A, fake_B and real_B, respectively. For other training arguments inherited from [pix2pix](https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix), check out the modules in `./options`.\n\n#### Train your own model\nSince ATME is originally introduced for supervised image-to-image translation, you will need to create an aligned dataset suitable for your application. Once you have done that, there is some freedom to select your best model, in particular, by tweeking the generator. \n\n\u003cp align=\"center\"\u003e\n  \u003cimg src='imgs/atme_backbone.png'\u003e\n\u003c/p\u003e\n\nThe UNet comprising the generator in ATME has a default configuration ([lucidrains](https://github.com/lucidrains/denoising-diffusion-pytorch/blob/main/README.md)) consisting of an embedding dimension $d=64$ and attention resolutions $R=(1,2,4,8)$. In general, $R=(r_1, r_2, \\cdots, r_H)$ represent the rates at which the feature maps grow during downsampling the input image through the UNet. That is, at the $i$-th layer, the number of feature maps go from $r_{i-1}d$ to $r_id$. You can vary the parameters $d$ and $R$ by using the arguments `--dim` and `--dim_mults`, respectively, in the training command above, e.g. `--dim 64 --dim_mults (1,2,4,8)` are the defaults. Such a variation is restricted by each $r_id$ being divisible by the number of groups `--groups` in the GroupNorm within each ResnetBlock, which defaults to `--groups 8`.\n\nThe supported arguments for the generator are\n```commandline\n--dim                        embedding dimension of the UNet (int)\n--dim_mults                  attention resolutions (int)\n--groups                     number of groups for GroupNorm within ResnetBlocks (int)\n--init_dim                   number of output channels after initial conv2d of x_t (int)\n--learned_sinusoidal_cond    learn fourier features for positional embedding? (bool)\n--random_fourier_features    random fourier features for positional embedding? (bool)\n--learned_sinusoidal_dim     twice the number of fourier frequencies to learn (int)\n--time_dim_mult              dim * time_dim_mult: number of output channels after time-MLP (int)\n```\nBy default, the positional embedding is carried out using fixed frequencies as in [Transformers](https://proceedings.neurips.cc/paper_files/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf).\n\n## Evaluation\nTo test your model, just run\n```\npython test.py --name PROJECT --model atme --direction AtoB --dataroot ./datasets/DATASET\n```\nThe results are placed in `./results/PROJECT`. Once there, you can compute the KID between real and generated images as\n```\nbash ./scripts/compute_kid.sh PROJECT\n```\n\n## Results\nBesides obtaining state-of-the-art results in supervised image-to-image translation, at a lesser cost than popular GAN and diffusion models, our main contribution is allowing a GAN to converge towards the theoretical optimum, as observed in the figure below, where the GAN objective in ATME converges towards $-\\log(4)$ on average. For discussion of the observed deviations take a look at the paper. \n\n\u003cimg src='imgs/convergence.png'\u003e\n\u003cbr\u003e\n\nYou can get the data for such plots from `./checkpoints/PROJECT/loss_log.csv`, by adding `D_real` and `D_fake`, as per Eqs. (1) and (2) in the paper.\n\n### Pre-trained models\nYou can find the original checkpoints, logs, and configurations used in the paper in the following links:\n\n| Dataset | Translation | Checkpoints |\n|-----| ------| -----|\n| Facades | A $\\rightarrow$ B | [[download]](https://zenodo.org/records/10065825/files/facades_AB.zip?download=1) |\n| Facades | B $\\rightarrow$ A | [[download]](https://zenodo.org/records/10065825/files/facades_BA.zip?download=1) |\n| Maps | A $\\rightarrow$ B | [[download]](https://zenodo.org/records/10065825/files/maps_AB.zip?download=1) |\n| Maps | B $\\rightarrow$ A | [[download]](https://zenodo.org/records/10065825/files/maps_BA.zip?download=1) |\n| Cityscapes | A $\\rightarrow$ B | [[download]](https://zenodo.org/records/10065825/files/cityscapes_AB.zip?download=1) |\n| Cityscapes | B $\\rightarrow$ A | [[download]](https://zenodo.org/records/10065825/files/cityscapes_BA.zip?download=1) |\n| Night2day | A $\\rightarrow$ B | [[download]](https://zenodo.org/records/10065825/files/night2day_AB.zip?download=1) |\n| Night2day | B $\\rightarrow$ A | [[download]](https://zenodo.org/records/10065825/files/night2day_BA.zip?download=1) |\n\n\nNote: the W models were originally named T. So make sure to call them W after download, if you want to load them. \n\n### Is $W_t$ a Brownian motion?\n\nThe script for training now has an optional argument `n_save_noisy` to allow studying the properties of $D_t$ and $W_t$. The animation below shows the tendency of the (rescaled) $D_t$ towards a flat distribution for all entries, corresponding to the Nash equilibrium.\nTake a look a the notebook `analyze_DW.ipynb` for experiments with the Maps dataset (in the direction AtoB) from which these results are obtained. You can find there some tests showing that $W_t$ behaves as a Brownian motion in a significant number of cases.\n\n\u003ctable\u003e\n  \u003ctr\u003e\n    \u003ctd style=\"text-align: center\"\u003e\u003cimg src=\"imgs/DWdW_452.gif\"\u003e\u003c/td\u003e\n    \u003ctd \u003e\u003cimg src=\"imgs/DWdW_536.gif\"\u003e\u003c/td\u003e\n  \u003c/tr\u003e\n  \u003ctr\u003e\n    \u003ctd style=\"text-align: center\"\u003e\u003cimg src=\"imgs/DWdW_690.gif\"\u003e\u003c/td\u003e\n    \u003ctd style=\"text-align: center\"\u003e\u003cimg src=\"imgs/DWdW_856.gif\"\u003e\u003c/td\u003e\n  \u003c/tr\u003e\n\u003c/table\u003e\n\n## Things left to do:\n- [x] Check if $W_t$ behaves as a Brownian motion.\n- [ ] Ablation of block (f) in the paper and how it leads to failure modes during training.\n- [ ] Ablation of block (g) in the paper and how it could generalize to different image resolutions.\n- [ ]  Generalization of ATME to unsupervised image-to-image translation.\n\n## Citation\n```\n@misc{solanocarrillo2023look,\n      title={Look ATME: The Discriminator Mean Entropy Needs Attention}, \n      author={Edgardo Solano-Carrillo and Angel Bueno Rodriguez and Borja Carrillo-Perez and Yannik Steiniger and Jannis Stoppe},\n      year={2023},\n      eprint={2304.09024},\n      archivePrefix={arXiv},\n      primaryClass={cs.CV}\n}\n\n@misc{solano_carrillo_2023_10065825,\n  author       = {Solano-Carrillo, Edgardo and\n                  Bueno Rodriguez, Angel and\n                  Carrillo-Perez, Borja and\n                  Steiniger, Yannik and\n                  Stoppe, Jannis},\n  title        = {ATME: Trained models and logs},\n  month        = nov,\n  year         = 2023,\n  publisher    = {Zenodo},\n  doi          = {10.5281/zenodo.10065825},\n  url          = {https://doi.org/10.5281/zenodo.10065825}\n}\n```","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fdlr-mi%2Fatme","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fdlr-mi%2Fatme","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fdlr-mi%2Fatme/lists"}