{"id":13737690,"url":"https://github.com/lucidrains/contrastive-learner","last_synced_at":"2025-05-06T23:46:13.343Z","repository":{"id":57417558,"uuid":"259756875","full_name":"lucidrains/contrastive-learner","owner":"lucidrains","description":"A simple to use pytorch wrapper for contrastive self-supervised learning on any neural network","archived":false,"fork":false,"pushed_at":"2021-03-12T16:36:03.000Z","size":320,"stargazers_count":133,"open_issues_count":2,"forks_count":13,"subscribers_count":6,"default_branch":"master","last_synced_at":"2025-03-31T04:11:47.326Z","etag":null,"topics":["artificial-intelligence","contrastive-loss","deep-learning","self-supervised-learning"],"latest_commit_sha":null,"homepage":null,"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/lucidrains.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}},"created_at":"2020-04-28T21:36:01.000Z","updated_at":"2025-03-20T22:28:03.000Z","dependencies_parsed_at":"2022-08-30T16:43:44.257Z","dependency_job_id":null,"html_url":"https://github.com/lucidrains/contrastive-learner","commit_stats":null,"previous_names":[],"tags_count":1,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fcontrastive-learner","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fcontrastive-learner/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fcontrastive-learner/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fcontrastive-learner/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/lucidrains","download_url":"https://codeload.github.com/lucidrains/contrastive-learner/tar.gz/refs/heads/master","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":252788406,"owners_count":21804280,"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","contrastive-loss","deep-learning","self-supervised-learning"],"created_at":"2024-08-03T03:01:57.388Z","updated_at":"2025-05-06T23:46:13.300Z","avatar_url":"https://github.com/lucidrains.png","language":"Python","funding_links":[],"categories":["Python"],"sub_categories":[],"readme":"\u003cimg src=\"./diagram.png\" width=\"450px\"/\u003e\n\n# Contrastive learning in Pytorch, made simple\n\n[![PyPI version](https://badge.fury.io/py/contrastive-learner.svg)](https://badge.fury.io/py/contrastive-learner)\n\nIt seems we have lift-off for self-supervised learning on images.\n\nThis is a simple to use Pytorch wrapper to enable contrastive self-supervised learning on any visual neural network. At the moment, it contains enough settings for one to train on either of the schemes used in SimCLR or CURL.\n\nYou can wrap any neural network that accepts a visual input, be it a resnet, policy network, or the discriminator of a GAN. The rest is taken care of.\n\n## Issues\n\nIt has surfaced that the results of CURL are \u003ca href=\"https://www.reddit.com/r/MachineLearning/comments/grnz0d/d_issues_reproducing_curl_algorithm_seems_broken/\"\u003enot reproducible\u003c/a\u003e. It is recommended that you go with the SimCLR settings until further notice.\n\n## Install\n\n```bash\n$ pip install contrastive-learner\n```\n\n## Usage\n\nSimCLR (projection head with normalized temperature-scaled cross-entropy loss)\n\n```python\nimport torch\nfrom contrastive_learner import ContrastiveLearner\nfrom torchvision import models\n\nresnet = models.resnet50(pretrained=True)\n\nlearner = ContrastiveLearner(\n    resnet,\n    image_size = 256,\n    hidden_layer = 'avgpool',  # layer name where output is hidden dimension. this can also be an integer specifying the index of the child\n    project_hidden = True,     # use projection head\n    project_dim = 128,         # projection head dimensions, 128 from paper\n    use_nt_xent_loss = True,   # the above mentioned loss, abbreviated\n    temperature = 0.1,         # temperature\n    augment_both = True        # augment both query and key\n)\n\nopt = torch.optim.Adam(learner.parameters(), lr=3e-4)\n\ndef sample_batch_images():\n    return torch.randn(20, 3, 256, 256)\n\nfor _ in range(100):\n    images = sample_batch_images()\n    loss = learner(images)\n    opt.zero_grad()\n    loss.backward()\n    opt.step()\n\n```\n\nCURL (with momentum averaged key encoder)\n\n```python\nimport torch\nfrom contrastive_learner import ContrastiveLearner\nfrom torchvision import models\n\nresnet = models.resnet50(pretrained=True)\n\nlearner = ContrastiveLearner(\n    resnet,\n    image_size = 256,\n    hidden_layer = 'avgpool',\n    use_momentum = True,         # use momentum for key encoder\n    momentum_value = 0.999,\n    project_hidden = False,      # no projection heads\n    use_bilinear = True,         # in paper, logits is bilinear product of query / key\n    use_nt_xent_loss = False,    # use regular contrastive loss\n    augment_both = False         # in curl, only the key is augmented\n)\n\nopt = torch.optim.Adam(learner.parameters(), lr=3e-4)\n\ndef sample_batch_images():\n    return torch.randn(20, 3, 256, 256)\n\nfor _ in range(100):\n    images = sample_batch_images()\n    loss = learner(images)\n    opt.zero_grad()\n    loss.backward()\n    opt.step()\n    learner.update_moving_average() # update moving average of key encoder\n```\n\n## Advanced\n\nIf you want to accumulate queries and keys to do contrastive loss on a bigger batch, use the `accumulate` keyword on the forward pass.\n\n```python\nfor _ in range(100):\n    for _ in range(5):\n        images = sample_batch_images()\n        _ = learner(images, accumulate=True)  # accumulate queries and keys\n    loss = learner.calculate_loss()           # calculate similarity on all accumulated\n    opt.zero_grad()\n    loss.backward()\n    opt.step()\n```\n\nBy default, this will use the augmentations recommended in the SimCLR paper, mainly color jitter, gaussian blur, and random resize crop. However, if you would like to specify your own augmentations, you can simply pass in a `augment_fn` in the constructor. Augmentations must work in the tensor space. If you decide to use torchvision augmentations, make sure the function converts first to PIL `.toPILImage()` and then back to tensors `.ToTensor()`\n\n```python\ncustom_augment_fn = nn.Sequential(\n    kornia.augmentations.RandomHorizontalFlip()\n)\n\nlearner = ContrastiveLearner(\n    resnet,\n    image_size = 256,\n    hidden_layer = -2,\n    project_hidden = True,\n    project_dim = 128,\n    use_nt_xent_loss = True,\n    augment_fn = custom_augment_fn\n)\n```\n\n## Citations\n\n```bibtex\n@misc{chen2020simple,\n    title   = {A Simple Framework for Contrastive Learning of Visual Representations},\n    author  = {Ting Chen and Simon Kornblith and Mohammad Norouzi and Geoffrey Hinton},\n    year    = {2020}\n}\n```\n\n```bibtex\n@misc{srinivas2020curl,\n    title   = {CURL: Contrastive Unsupervised Representations for Reinforcement Learning},\n    author  = {Aravind Srinivas and Michael Laskin and Pieter Abbeel},\n    year    = {2020}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Fcontrastive-learner","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Flucidrains%2Fcontrastive-learner","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Fcontrastive-learner/lists"}