{"id":18478809,"url":"https://github.com/facebookresearch/ssl-data-curation","last_synced_at":"2025-04-12T04:53:36.117Z","repository":{"id":241339694,"uuid":"805292666","full_name":"facebookresearch/ssl-data-curation","owner":"facebookresearch","description":"PyTorch code for hierarchical k-means -- a data curation method for self-supervised learning","archived":false,"fork":false,"pushed_at":"2024-06-21T08:54:19.000Z","size":2011,"stargazers_count":151,"open_issues_count":2,"forks_count":11,"subscribers_count":27,"default_branch":"main","last_synced_at":"2025-04-12T04:53:26.690Z","etag":null,"topics":[],"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/facebookresearch.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":"CONTRIBUTING.md","funding":null,"license":"LICENSE","code_of_conduct":"CODE_OF_CONDUCT.md","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":"2024-05-24T09:18:30.000Z","updated_at":"2025-04-07T13:53:03.000Z","dependencies_parsed_at":"2024-05-27T20:54:14.677Z","dependency_job_id":"1c3760f8-d62f-413f-acba-b97b0b7c86f6","html_url":"https://github.com/facebookresearch/ssl-data-curation","commit_stats":null,"previous_names":["facebookresearch/ssl-data-curation"],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/facebookresearch%2Fssl-data-curation","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/facebookresearch%2Fssl-data-curation/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/facebookresearch%2Fssl-data-curation/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/facebookresearch%2Fssl-data-curation/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/facebookresearch","download_url":"https://codeload.github.com/facebookresearch/ssl-data-curation/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":248519472,"owners_count":21117757,"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":[],"created_at":"2024-11-06T12:12:02.287Z","updated_at":"2025-04-12T04:53:36.097Z","avatar_url":"https://github.com/facebookresearch.png","language":"Jupyter Notebook","funding_links":[],"categories":["public repositories"],"sub_categories":[],"readme":"# Automatic Data Curation for Self-Supervised Learning: A Clustering-Based Approach\n\n**[FAIR at Meta](https://ai.facebook.com/research/)**\n\n*Huy V. Vo,\nVasil Khalidov,\nTimothée Darcet,\nThéo Moutakanni,\nNikita Smetanin,\nMarc Szafraniec,\nHugo Touvron,\nCamille Couprie,\nMaxime Oquab,\nArmand Joulin,\nHervé Jégou,\nPatrick Labatut,\nPiotr Bojanowski*\n\nPyTorch implementation for the data curation pipeline with hierarchical k-means. For more detail, see the paper **[Automatic Data Curation for Self-Supervised Learning: A Clustering-Based Approach](https://arxiv.org/abs/2405.15613)**.\n\n\u003cp align=\"center\"\u003e\n  \u003cimg width=\"95%\" alt=\"data curation pipeline\" src=\"images/curation_pipeline.png\"\u003e\n\u003c/p\u003e \n\n## Contents\n- [Installation](#installation)\n- [Running hierarchical k-means](#running-hierarchical-k-means)\n  * [On small data](#on-small-data)\n  * [On large data](#on-large-data)\n- [Notebook](#notebook)\n- [Contributing](#contributing)\n- [License](#license)\n- [Citation](#citation)\n\n## Installation\n```\ngit clone git@github.com:facebookresearch/ssl-data-curation.git\ncd ssl-data-curation\nconda create -n ssl-data-curation python=3.10\nconda activate ssl-data-curation\npip install -r requirements.txt\n```\n\n## Running hierarchical k-means\n### On small data\nWe provide below an example of a 2-level hierarchical k-means on a small toy random dataset. We first run hierarchical k-means on the toy dataset then sample 1000 points from it with hierarchical sampling. A visualisation is provided in [vis/notebook.ipynb](vis/notebook.ipynb).\n```\nimport torch\nimport numpy as np\n\nfrom src.clusters import HierarchicalCluster\nfrom src import (\n  hierarchical_kmeans_gpu as hkmg,\n  hierarchical_sampling as hs\n)\n\ndef make_ring(n, rmin, rmax):\n    r = np.random.rand(n) * (rmax - rmin) + rmin\n    alpha = np.random.rand(n) * 2 * np.pi\n    return np.vstack([r * np.cos(alpha), r * np.sin(alpha)]).T\n\ndata = np.concatenate([\n    make_ring(20000, 0.7, 1.0) + np.array([-2.2, 1.]),\n    make_ring(200, 0.7, 1.0) + np.array([0., 1.]),\n    make_ring(1000, 0.7, 1.0) + np.array([2.2, 1.]),\n    make_ring(500, 0.7, 1.0) + np.array([-1.2, 0.2]),\n    make_ring(8000, 0.7, 1.0) + np.array([1.2, 0.2]),\n])\n\nclusters = hkmg.hierarchical_kmeans_with_resampling(\n  data=torch.tensor(data, device=\"cuda\", dtype=torch.float32),\n  n_clusters=[1000, 300],\n  n_levels=2,\n  sample_sizes=[15, 2],\n  verbose=False,\n)\n\ncl = HierarchicalCluster.from_dict(clusters)\nsampled_indices = hs.hierarchical_sampling(cl, target_size=1000)\n```\n\n\u003cp align=\"center\"\u003e\n  \u003cimg width=\"95%\" alt=\"data curation pipeline\" src=\"images/toy_example.png\"\u003e\n\u003c/p\u003e \n\n### On large data\nTo launch hierarchical k-means on large data, we need to prepare a config file. We provide below an example illustrating how to launch a 2-level hierarchical k-means on random embeddings with config in [configs/2levels_random_embeddings.yaml](configs/2levels_random_embeddings.yaml).\n```\n# Prepare the experiment\ncd ssl-data-curation\nmkdir -p data\ncd scripts\npython -c 'import numpy as np; np.save( \"../data/100k_random.npy\", np.random.randn(100000,256))'\npython hierarchical_kmeans_launcher.py \\\n  --exp_dir ../data/2levels_random_embeddings \\\n  --embeddings_path ../data/100k_random.npy \\\n  --config_file ../configs/2levels_random_embeddings.yaml\n\ncd ../data/2levels_random_embeddings\n# Launch with slurm\nbash launcher.sh\n# Launch locally if only 1 node is used\n# bash local_launcher.sh\n\ncd ssl-data-curation/scripts\n# Sampled indices will be saved in ssl-data-curation/data/2levels_random_embeddings/curated_datasets\nPYTHONPATH=.. python run_hierarchical_sampling.py \\\n  --clustering_path ../data/2levels_random_embeddings \\\n  --target_size 20000 \\\n  --save\n```\n\nWe also  provide the config used for our web-based image data pool in [configs/4levels_web_based_images.yaml](configs/4levels_web_based_images.yaml). \n\n## Notebook\nWe provide a [notebook](vis/notebook.ipynb) to reproduce visualizations in the paper and show additional examples. \n\n## Contributing\nSee [contributing](CONTRIBUTING.md) and the [code of conduct](CODE_OF_CONDUCT.md).\n\n## License\nThis code is CC-BY-NC 4.0 licensed, as found in [LICENSE](LICENSE).\n\n## Citation\nIf you find our work useful, please consider giving a star and a citation:\n```\n@article{vo2024automatic,\n  title={Automatic Data Curation for Self-Supervised Learning: A Clustering-Based Approach},\n  author={Vo, Huy V. and Khalidov, Vasil and Darcet, Timoth{\\'e}e and Moutakanni, Th{\\'e}o and Smetanin, Nikita and Szafraniec, Marc and Touvron, Hugo and Couprie, Camille and Oquab, Maxime and Joulin, Armand and Jégou, Hervé and Labatut, Patrick and Bojanowski, Piotr},\n  journal={arXiv:2405.15613},\n  year={2024},\n}\n```","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Ffacebookresearch%2Fssl-data-curation","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Ffacebookresearch%2Fssl-data-curation","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Ffacebookresearch%2Fssl-data-curation/lists"}