{"id":28753492,"url":"https://github.com/google-deepmind/distribution_shift_framework","last_synced_at":"2025-06-17T00:40:24.661Z","repository":{"id":66141111,"uuid":"470912097","full_name":"google-deepmind/distribution_shift_framework","owner":"google-deepmind","description":"This repository contains the code of the distribution shift framework presented in A Fine-Grained Analysis on Distribution Shift (Wiles et al., 2022).","archived":false,"fork":false,"pushed_at":"2024-05-06T17:28:33.000Z","size":68,"stargazers_count":77,"open_issues_count":6,"forks_count":7,"subscribers_count":7,"default_branch":"master","last_synced_at":"2024-05-06T18:45:22.367Z","etag":null,"topics":["artificial-intelligence","deep-learning","distribution-shift","jax","machine-learning","neural-networks"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/google-deepmind.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":"CONTRIBUTING.md","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":"2022-03-17T08:43:06.000Z","updated_at":"2024-04-13T19:46:52.000Z","dependencies_parsed_at":"2023-09-07T20:34:50.789Z","dependency_job_id":"2eb68c6c-a6ce-43d2-ad5c-9f7a0e6c54d9","html_url":"https://github.com/google-deepmind/distribution_shift_framework","commit_stats":null,"previous_names":["google-deepmind/distribution_shift_framework"],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/google-deepmind/distribution_shift_framework","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google-deepmind%2Fdistribution_shift_framework","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google-deepmind%2Fdistribution_shift_framework/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google-deepmind%2Fdistribution_shift_framework/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google-deepmind%2Fdistribution_shift_framework/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/google-deepmind","download_url":"https://codeload.github.com/google-deepmind/distribution_shift_framework/tar.gz/refs/heads/master","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google-deepmind%2Fdistribution_shift_framework/sbom","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":260268635,"owners_count":22983601,"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":["artificial-intelligence","deep-learning","distribution-shift","jax","machine-learning","neural-networks"],"created_at":"2025-06-17T00:40:23.820Z","updated_at":"2025-06-17T00:40:24.627Z","avatar_url":"https://github.com/google-deepmind.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# Distribution Shift Framework\n\nThis repository contains the code of the distribution shift framework presented\nin [A Fine-Grained Analysis on Distribution Shift](https://openreview.net/forum?id=Dl4LetuLdyK)\n(Wiles et al., 2022).\n\n## Contents\n\nThe framework allows to train models with different training methods on\ndatasets undergoing specific kinds of distribution shift.\n\n### Training Methods\n\nCurrently the following training methods are supported (by setting the\n`algorithm` [config option](#config-options)):\n\n* Empirical Risk Minimization (**ERM**, [Vapnik, 1992](https://papers.nips.cc/paper/1991/hash/ff4d5fbbafdf976cfdc032e3bde78de5-Abstract.html))\n* Invariant Risk Minimization (**IRM**, [Arjovsky et al., 2019](https://arxiv.org/abs/1907.02893))\n* Deep Correlation Alignment (**Deep CORAL**, [Sun \u0026 Saenko, 2016](https://link.springer.com/chapter/10.1007/978-3-319-49409-8_35))\n* Domain-Adversarial Training of Neural Networks (**DANN**, [Ganin et al., 2016](https://jmlr.org/papers/v17/15-239.html))\n* Style-Agnostic Networks (**SagNet**, [Nam et al., 2021](https://openaccess.thecvf.com/content/CVPR2021/html/Nam_Reducing_Domain_Gap_by_Reducing_Style_Bias_CVPR_2021_paper.html))\n* (Batch Normalization Adaption (**BN-Adapt**, [Schneider et al., 2020](https://proceedings.neurips.cc/paper/2020/hash/85690f81aadc1749175c187784afc9ee-Abstract.html))\n* Just Train Twice (**JTT**, [Liu et al., 2021](http://proceedings.mlr.press/v139/liu21f.html))\n* Inter-domain Mixup (**MixUp**, [Gulrajani \u0026 LopezPaz, 2021](https://openreview.net/forum?id=lQdXeXDoWtI))\n\n### Model Architectures\n\nThe `model` [config option](#config-options) can be set to one of the following\narchitectures\n\n* ResNet18, ResNet50, ResNet101 ([He et al., 2016](https://ieeexplore.ieee.org/document/7780459))\n* MLP ([Vapnik, 1992](https://papers.nips.cc/paper/1991/hash/ff4d5fbbafdf976cfdc032e3bde78de5-Abstract.html))\n\n### Datasets\n\nYou can train on the following datasets (by setting the `dataset_name`\n[config option.](#config-options)):\n\n* dSprites ([Matthey et al., 2017](https://github.com/deepmind/dsprites-dataset))\n* SmallNorb ([LeCun et al., 2004](https://ieeexplore.ieee.org/document/1315150))\n* Shapes3D ([Burgess \u0026 Kim, 2018](https://github.com/deepmind/3d-shapes))\n\nEach dataset has a task (e.g. shape prediction on dSprites, set with the `label`\n[config option](#config-options)) and a set of properties (e.g. the colour of\nthe shape in dSprites, set with the `property_label`\n[config option](#config-options)).\n\n### Distribution Shift Scenarios\n\nYou can evaluate your model on different conditions by varying the distribution\nof labels and properties in the configs. For each part of the distribution,\nyou then assign a probability of sampling from that part of the distribution.\n\n* **Unseen data shift** (`ood`): Some parts of the distribution of the property\n    are unseen at training time (e.g. certain colours may be unseen in\n    dSprites).\n* **Spurious correlation** (`correlated`): Some property is correlated with the\n    label at training time but not at test (e.g. all circles are red in\n    training).\n* **Low data drift** (`lowdata`): Certain combinations of label and property are seen at a\n    a lower rate during training while they are uniformly distributed during\n    test.\n\nAdditionally you can modify these scenarios with two conditions:\n\n* **Label noise** (`noise`): A certain percentage of the training labels are\n    corrupted.\n* **Fixed dataset size** (`fixeddata`): We reduce the total training dataset\n    size to a fixed amount.\n\nThese scenarios can be set through the `test_case`\n[config option.](#config-options)) with the keywords in parenthesis and an\noptional modifier separated by a full stop, e.g. `lowdata.noise` for low data\ndrift with added label noise.\n\n### Future Additions\n\nWe plan to add additional methods, models and datasets from the paper as well\nas the raw results from all the experiments.\n\n## Usage Instructions\n\n### Installing\n\nThe following has been tested using Python 3.9.9.\n\nFor GPU support with JAX, edit `requirements.txt` before running `run.sh`\n(e.g., use `jaxline==0.1.67+cuda111`). See JAX's installation\n[instructions](https://github.com/google/jax#installation) for more details.\n\nExecute `run.sh` to create and activate a virtualenv, install all necessary\ndependencies and run a test program to ensure that you can import all the\nmodules.\n\n```\n# Run from the parent directory.\nsh distribution_shift_framework/run.sh\n```\n\n\n### Running the Code\n\nTo train a model, use this virtualenv:\n\n```\nsource /tmp/distribution_shift_framework/bin/activate\n```\n\nand then run\n\n```\npython3 -m distribution_shift_framework.classification.experiment \\\n--jaxline_mode=train \\\n--config=distribution_shift_framework/classification/config.py\n```\n\nFor evaluation run\n\n```\npython3 -m distribution_shift_framework.classification.experiment \\\n--jaxline_mode=eval \\\n--config=distribution_shift_framework/classification/config.py\n```\n\n### Config Options {#config-options}\n\nCommon changes can be done through an options string following the config file.\nThe following options are available:\n\n* `algorithm`: What training method to use for training.\n* `model`:: The model architecture to evaluate.\n* `dataset_name`: The name of the dataset.\n* `test_case`: Which of the distribution shift scenarios to set up.\n* `label`: The label we're predicting.\n* `property_label`: Which property is treated as in or out of\n  distribution (for the ood test_case), is correlated with the label\n  (for the correlated setup) and is treated as having a low data region\n  (for the low_data setup).\n* `number_of_seeds`: How many seeds to sweep over.\n* `batch_size`: Batch size used for training and evaluation.\n* `training_steps`: How many steps to train for.\n* `pretrained_checkpoint`: Path to a checkpoint for a pretrained model.\n* `overwrite_image_size`: Height and width to resize the images to. 0 means\n  no resizing.\n* `eval_specific_ckpt`: Path to a checkpoint for a one time evaluation.\n* `wids`: Which wids of the checkpoint to look at.\n* `sweep_index`: Which experiment from the sweep to run.\n* `use_fake_data`: Whether to use fake data for testing.\n\n\nMultiple options need to be separated by commas. An example would be\n\n```\npython3 -m distribution_shift_framework.classification.experiment \\\n--jaxline_mode=train \\\n--config=distribution_shift_framework/classification/config.py:algorithm=SagNet,test_case=lowdata.noise,model=truncatedresnet18,property_label=label_object_hue,label=label_shape,dataset_name=shapes3d\n```\n\nWhich would train a **truncated ResNet18** with the **SagNet** algorithm in the\n**low data** setting with added **label noise** on the **Shapes3D** dataset.\n**Shape** is used as the label for classification while **object hue** is used\nas the property that the distribution shifts over.\n\n### Sweeps\n\nBy default the program generates sweeps over multiple hyper-parameters depending\non the chosen training method, dataset and distribution shift scenario. The\n`sweep_index` option lets you choose which of the configs in the sweep you want\nto run.\n\n## Citing this work\n\nIf you use this code (or any derived code) in your work, please cite the\naccompanying paper:\n\n```\n@inproceedings{wiles2022fine,\n  title={A Fine-Grained Analysis on Distribution Shift},\n  author={Olivia Wiles and Sven Gowal and Florian Stimberg and Sylvestre-Alvise Rebuffi and Ira Ktena and Krishnamurthy Dj Dvijotham and Ali Taylan Cemgil},\n  booktitle={International Conference on Learning Representations},\n  year={2022},\n  url={https://openreview.net/forum?id=Dl4LetuLdyK}\n}\n```\n\n## License and Disclaimer\n\nCopyright 2022 DeepMind Technologies Limited.\n\nAll software is licensed under the Apache License, Version 2.0 (Apache 2.0);\nyou may not use this file except in compliance with the License. You may obtain\na copy of the Apache 2.0 license at\n\n[https://www.apache.org/licenses/LICENSE-2.0](https://www.apache.org/licenses/LICENSE-2.0)\n\nAll non-code materials are licensed under the Creative Commons Attribution 4.0\nInternational License (CC-BY License). You may obtain a copy of the CC-BY\nLicense at:\n\n[https://creativecommons.org/licenses/by/4.0/legalcode](https://creativecommons.org/licenses/by/4.0/legalcode)\n\nYou may not use the non-code portions of this file except in compliance with the\nCC-BY License.\n\nUnless required by applicable law or agreed to in writing, software distributed\nunder the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR\nCONDITIONS OF ANY KIND, either express or implied. See the License for the\nspecific language governing permissions and limitations under the License.\n\nThis is not an official Google product.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fgoogle-deepmind%2Fdistribution_shift_framework","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fgoogle-deepmind%2Fdistribution_shift_framework","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fgoogle-deepmind%2Fdistribution_shift_framework/lists"}