Package: open-clip-torch-any-py3
OpenCLIP[Paper][Colab]Welcome to an open source implementation of OpenAI'sCLIP(Contrastive Language-Image Pre-training).The goal of this repository is to enable training models with contrastive image-text supervision, and to investigate their properties such as robustness to distribution shift. Our starting point is an implementation of CLIP that matches the accuracy of the original CLIP models when trained on the same dataset.
Specifically, a ResNet-50 model trained with our codebase on OpenAI's15 million image subset of YFCCachieves32.7%top-1 accuracy on ImageNet. OpenAI's CLIP model reaches31.3%when trained on the same subset of YFCC. For ease of experimentation, we also provide code for training on the 3 million images in theConceptual Captionsdataset, where a ResNet-50x4 trained with our codebase reaches 22.2% top-1 ImageNet accuracy.We further this with a replication study on a dataset of comparable size to OpenAI's. UsingLAION-400M, we train CLIP with aViT-B/32 and achieve an accuracy of62.9%, comparable to OpenAI's63.2%, zero-shot top-1 on ImageNet1kViT-B/16 and achieve an accuracy of67.1%, comparable to OpenAI's68.3%(as measured here, 68.6% in paper)ViT-B/16+ 240x240 (~50% more FLOPS than B/16 224x224) and achieve an accuracy of69.2%ViT-L/14 and achieve an accuracy of72.77%, vs OpenAI's75.5%(as measured here, 75.3% in paper)As we describe in more detailbelow, CLIP models in a medium accuracy regime already allow us to draw conclusions about the robustness of larger CLIP models since the models followreliable scaling laws.This codebase is work in progress, and we invite all to contribute in making it more acessible and useful. In the future, we plan to add support for TPU training and release larger models. We hope this codebase facilitates and promotes further research in contrastive image-text learning. Please submit an issue or send an email if you have any other requests or suggestions.Note that portions ofsrc/open_clip/modelling and tokenizer code are adaptations of OpenAI's officialrepository.ApproachImage Credit:https://github.com/openai/CLIPUsagepip install open_clip_torchimporttorchfromPILimportImageimportopen_clipmodel,_,preprocess=open_clip.create_model_and_transforms('ViT-B-32-quickgelu',pretrained='laion400m_e32')image=preprocess(Image.open("CLIP.png")).unsqueeze(0)text=open_clip.tokenize(["a diagram","a dog","a cat"])withtorch.no_grad():image_features=model.encode_image(image)text_features=model.encode_text(text)image_features/=image_features.norm(dim=-1,keepdim=True)text_features/=text_features.norm(dim=-1,keepdim=True)text_probs=(100.0*image_features@text_features.T).softmax(dim=-1)print("Label probs:",text_probs)# prints: [[1., 0., 0.]]To compute billions of embeddings efficiently, you can useclip-retrievalwhich has openclip support.Fine-tuning on classification tasksThis repository is focused on training CLIP models. To fine-tune atrainedzero-shot model on a downstream classification task such as ImageNet, please seeour other repository: WiSE-FT. TheWiSE-FT repositorycontains code for our paper onRobust Fine-tuning of Zero-shot Models, in which we introduce a technique for fine-tuning zero-shot models while preserving robustness under distribution shift.DataConceptual CaptionsOpenCLIP reads a CSV file with two columns: a path to an image, and a text caption. The names of the columns are passed as an argument tomain.py.The scriptsrc/data/gather_cc.pywill collect the Conceptual Captions images. First, download theConceptual Captions URLsand then run the script from our repository:python3src/data/gather_cc.pypath/to/Train_GCC-training.tsvpath/to/Validation_GCC-1.1.0-Validation.tsvOur training set contains 2.89M images, and our validation set contains 13K images.YFCC and other datasetsIn addition to specifying the training data via CSV files as mentioned above, our codebase also supportswebdataset, which is recommended for larger scale datasets. The expected format is a series of.tarfiles. Each of these.tarfiles should contain two files for each training example, one for the image and one for the corresponding text. Both files should have the same name but different extensions. For instance,shard_001.tarcould contain files such asabc.jpgandabc.txt. You can learn more aboutwebdatasetathttps://github.com/webdataset/webdataset. We use.tarfiles with 1,000 data points each, which we create usingtarp.You can download the YFCC dataset fromMultimedia Commons.
Similar to OpenAI, we used a subset of YFCC to reach the aforementioned accuracy numbers.
The indices of images in this subset are inOpenAI's CLIP repository.Training CLIPSetup Environment and Install dependenciesConda# Create a conda environment (heavily recommended)condacreate-nopen_clippython=3.10
condaactivateopen_clipInstall conda PyTorch as perhttps://pytorch.org/get-started/locally/Virtualenvopenclip also can be used with virtualenv with these lines:python3 -m venv .env
source .env/bin/activate
pip install -U pip
make installInstall pip PyTorch as perhttps://pytorch.org/get-started/locally/Test can be run withmake install-devthenmake testOther dependenciesInstall open_clip pacakge and remaining dependencies:cdopen_clip
pythonsetup.pyinstallIf you want to train models, you will also need to install the packages
fromrequirements-training.txt.Sample single-process running code:python-mtraining.main\--save-frequency1\--zeroshot-frequency1\--report-totensorboard\--train-data="/path/to/train_data.csv"\--val-data="/path/to/validation_data.csv"\--csv-img-keyfilepath\--csv-caption-keytitle\--imagenet-val=/path/to/imagenet/root/val/\--warmup10000\--batch-size=128\--lr=1e-3\--wd=0.1\--epochs=30\--workers=8\--modelRN50Note:imagenet-valis the path to thevalidationset of ImageNet for zero-shot evaluation, not the training set!
You can remove this argument if you do not want to perform zero-shot evaluation on ImageNet throughout training. Note that thevalfolder should contain subfolders. If it doest not, please usethis script.Multi-GPU and BeyondThis code has been battle tested up to 1024 A100s and offers a variety of solutions
for distributed training. We include native support for SLURM clusters.As the number of devices used to train increases, so does the space complexity of
the the logit matrix. Using a naïve all-gather scheme, space complexity will beO(n^2). Instead, complexity may become effectively linear if the flags--gather-with-gradand--local-lossare used. This alteration results in one-to-one
numerical results as the naïve method.Single-NodeWe make use oftorchrunto launch distributed jobs. The following launches a
a job on a node of 4 GPUs:cdopen_clip/src
torchrun--nproc_per_node4-mtraining.main\--train-data'/data/cc12m/cc12m-train-{0000..2175}.tar'\--train-num-samples10968539\--dataset-typewebdataset\--batch-size320\--precisionamp\--workers4\--imagenet-val/data/imagenet/validation/Multi-NodeThe same script above works, so long as users include information about the number
of nodes and host node.cdopen_clip/src
torchrun--nproc_per_node=4\--rdzv_endpoint=$HOSTE_NODE_ADDR\-mtraining.main\--train-data'/data/cc12m/cc12m-train-{0000..2175}.tar'\--train-num-samples10968539\--dataset-typewebdataset\--batch-size320\--precisionamp\--workers4\--imagenet-val/data/imagenet/validation/SLURMThis is likely the easiest solution to utilize. The following script was used to
train our largest models:#!/bin/bash -x#SBATCH --nodes=32#SBATCH --gres=gpu:4#SBATCH --ntasks-per-node=4#SBATCH --cpus-per-task=6#SBATCH --wait-all-nodes=1#SBATCH --job-name=open_clip#SBATCH --account=ACCOUNT_NAME#SBATCH --partition PARTITION_NAMEeval"$(/path/to/conda/bin/condashell.bashhook)"# init condacondaactivateopen_clipexportCUDA_VISIBLE_DEVICES=0,1,2,3exportMASTER_PORT=12802master_addr=$(scontrolshowhostnames"$SLURM_JOB_NODELIST"|head-n1)exportMASTER_ADDR=$master_addrcd/shared/open_clipexportPYTHONPATH="$PYTHONPATH:$PWD/src"srun--cpu_bind=v--accel-bind=gnpython-usrc/training/main.py\--save-frequency1\--report-totensorboard\--train-data="/data/LAION-400M/{00000..41455}.tar"\--warmup2000\--batch-size=256\--epochs=32\--workers=8\--modelViT-B-32\--name"ViT-B-32-Vanilla"\--seed0\--local-loss\--gather-with-gradResuming from a checkpoint:python-mtraining.main\--train-data="/path/to/train_data.csv"\--val-data="/path/to/validation_data.csv"\--resume/path/to/checkpoints/epoch_K.ptLoss CurvesWhen run on a machine with 8 GPUs the command should produce the following training curve for Conceptual Captions:More detailed curves for Conceptual Captions are given at/docs/clip_conceptual_captions.md.When training a RN50 on YFCC the same hyperparameters as above are used, with the exception oflr=5e-4andepochs=32.Note that to use another model, likeViT-B/32orRN50x4orRN50x16orViT-B/16, specify with--model RN50x4.Launch tensorboard:tensorboard--logdir=logs/tensorboard/--port=7777Evaluation / Zero-ShotEvaluating local checkpoint:python-mtraining.main\--val-data="/path/to/validation_data.csv"\--modelRN101\--pretrained/path/to/checkpoints/epoch_K.ptEvaluating hosted pretrained checkpoint on ImageNet zero-shot prediction:python-mtraining.main\--imagenet-val/path/to/imagenet/validation\--modelViT-B-32-quickgelu\--pretrainedlaion400m_e32Pretrained model detailsLAION-400M -https://laion.ai/laion-400-open-datasetWe are working on reproducing OpenAI's ViT results with the comparably sized (and open) LAION-400M dataset. Trained
weights may be found in releasev0.2.The LAION400M weights have been trained on the JUWELS supercomputer (see acknowledgements section below).ViT-B/32 224x224We replicate OpenAI's results on ViT-B/32, reaching a top-1 ImageNet-1k zero-shot accuracy of 62.96%.Zero-shot comparison (courtesy of Andreas Fürst)ViT-B/32 was trained with 128 A100 (40 GB) GPUs for ~36 hours, 4600 GPU-hours. The per-GPU batch size was 256 for a global batch size of 32768. 256 is much lower than it could have been (~320-384) due to being sized initially before moving to 'local' contrastive loss.ViT-B/16 224x224The B/16 LAION400M training reached a top-1 ImageNet-1k zero-shot validation score of 67.07.This was the first major train session using the updated webdataset 0.2.x code. A bug was found that prevented shards from being shuffled properly between nodes/workers each epoch. This was fixed part way through training (epoch 26) but likely had an impact.ViT-B/16 was trained with 176 A100 (40 GB) GPUS for ~61 hours, 10700 GPU-hours. Batch size per GPU was 192 for a global batch size of 33792.ViT-B/16+ 240x240The B/16+ 240x240 LAION400M training reached a top-1 ImageNet-1k zero-shot validation score of 69.21.This model is the same depth as the B/16, but increases thevision width from 768 -> 896text width from 512 -> 640the resolution 224x224 -> 240x240 (196 -> 225 tokens)Unlike the B/16 run above, this model was a clean run with no dataset shuffling issues.ViT-B/16+ was trained with 224 A100 (40 GB) GPUS for ~61 hours, 13620 GPU-hours. Batch size per GPU was 160 for a global batch size of 35840.ViT-L/14 224x224The L/14 LAION-400M training reached a top-1 ImageNet-1k zero-shot validation score of 72.77.ViT-L/14 was trained with 400 A100 (40 GB) GPUS for ~127 hours, 50800 GPU-hours. Batch size per GPU was 96 for a global batch size of 38400. Grad checkpointing was enabled.LAION-2B (en) -https://laion.ai/laion-5b-a-new-era-of-open-large-scale-multi-modal-datasets/A ~2B sample subset of LAION-5B with english captions (https://huggingface.co/datasets/laion/laion2B-en)ViT-B/32 224x224A ViT-B/32 trained on LAION-2B, reaching a top-1 ImageNet-1k zero-shot accuracy of 65.62%.ViT-B/32 was trained with 112 A100 (40 GB) GPUs. The per-GPU batch size was 416 for a global batch size of 46592. Compute generously provided bystability.ai.YFCC-15MBelow are checkpoints of models trained on YFCC-15M, along with their zero-shot top-1 accuracies on ImageNet and ImageNetV2. These models were trained using 8 GPUs and the same hyperparameters described in the "Sample running code" section, with the exception oflr=5e-4andepochs=32.ResNet-50(32.7% / 27.9%)ResNet-101(34.8% / 30.0%)CC12M -https://github.com/google-research-datasets/conceptual-12mResNet-50(36.45%)Pretrained Model InterfaceWe offer a simple model interface to instantiate both pre-trained and untrained models.NOTE: Many existing checkpoints use the QuickGELU activation from the original OpenAI models. This activation is actually less efficient that native torch.nn.GELU in recent versions of PyTorch. The model defaults are now nn.GELU, so one should use model definitions with-quickgelupostfix for the OpenCLIP pretrained weights. All OpenAI pretrained weights will always default to QuickGELU. One can also use the non-quickgelumodel definitions with pretrained weights using QuickGELU but there will be an accuracy drop, for fine-tune that will likely vanish for longer runs.Future trained models will use nn.GELU.>>>importopen_clip>>>open_clip.list_pretrained()[('RN50','openai'),('RN50','yfcc15m'),('RN50','cc12m'),('RN50-quickgelu','openai'),('RN50-quickgelu','yfcc15m'),('RN50-quickgelu','cc12m'),('RN101','openai'),('RN101','yfcc15m'),('RN101-quickgelu','openai'),('RN101-quickgelu','yfcc15m'),('RN50x4','openai'),('RN50x16','openai'),('RN50x64','openai'),('ViT-B-32','openai'),('ViT-B-32','laion2b_e16'),('ViT-B-32','laion400m_e31'),('ViT-B-32','laion400m_e32'),('ViT-B-32-quickgelu','openai'),('ViT-B-32-quickgelu','laion400m_e31'),('ViT-B-32-quickgelu','laion400m_e32'),('ViT-B-16','openai'),('ViT-B-16','laion400m_e31'),('ViT-B-16','laion400m_e32'),('ViT-B-16-plus-240','laion400m_e31'),('ViT-B-16-plus-240','laion400m_e32'),('ViT-L-14','openai'),('ViT-L-14-336','openai')]>>>model,train_transform,eval_transform=open_clip.create_model_and_transforms('ViT-B-32',pretrained='laion2b_e16')Scaling trendsThe plot below shows how zero-shot performance of CLIP models varies as we scale the number of samples used for training. Zero-shot performance increases steadily for both ImageNet andImageNetV2, and is far from saturated at ~15M samples.Why are low-accuracy CLIP models interesting?TL;DR:CLIP models have high effective robustness, even at small scales.CLIP models are particularly intriguing because they are more robust to natural distribution shifts (see Section 3.3 in theCLIP paper).
This phenomena is illustrated by the figure below, with ImageNet accuracy on the x-axis
andImageNetV2(a reproduction of the ImageNet validation set with distribution shift) accuracy on the y-axis.
Standard training denotes training on the ImageNet train set and the CLIP zero-shot models
are shown as stars.As observed byTaori et al., 2020andMiller et al., 2021, the in-distribution
and out-of-distribution accuracies of models trained on ImageNet follow a predictable linear trend (the red line in the above plot).Effective robustnessquantifies robustness as accuracy beyond this baseline, i.e., how far a model lies above the red line. Ideally a model would not suffer from distribution shift and fall on the y = x line (trained human labelers are within a percentage point of the y = x line).Even though the CLIP models trained with
this codebase achieve much lower accuracy than those trained by OpenAI, our models still lie on the same
trend of improved effective robustness (the purple line). Therefore, we can study what makes
CLIP robust without requiring industrial-scale compute.For more information on effective robustness, please see:Recht et al., 2019.Taori et al., 2020.Miller et al., 2021.To know more about the factors that contribute to CLIP's robustness refer toFang et al., 2022.AcknowledgmentsWe gratefully acknowledge the Gauss Centre for Supercomputing e.V. (www.gauss-centre.eu) for funding this part of work by providing computing time through the John von Neumann Institute for Computing (NIC) on the GCS Supercomputer JUWELS Booster at Jülich Supercomputing Centre (JSC).The TeamCurrent development of this repository is led byRoss Wightman,Cade Gordon, andVaishaal Shankar.The original version of this repository is from a group of researchers at UW, Google, Stanford, Amazon, Columbia, and Berkeley.Gabriel Ilharco*,Mitchell Wortsman*,Nicholas Carlini,Rohan Taori,Achal Dave,Vaishaal Shankar,John Miller,Hongseok Namkoong,Hannaneh Hajishirzi,Ali Farhadi,Ludwig SchmidtSpecial thanks toJong Wook KimandAlec Radfordfor help with reproducing CLIP!CitingIf you found this repository useful, please consider citing:@software{ilharco_gabriel_2021_5143773,author={Ilharco, Gabriel andWortsman, Mitchell andWightman, Ross andGordon, Cade andCarlini, Nicholas andTaori, Rohan andDave, Achal andShankar, Vaishaal andNamkoong, Hongseok andMiller, John andHajishirzi, Hannaneh andFarhadi, Ali andSchmidt, Ludwig},title={OpenCLIP},month=jul,year=2021,note={If you use this software, please cite it as below.},publisher={Zenodo},version={0.1},doi={10.5281/zenodo.5143773},url={https://doi.org/10.5281/zenodo.5143773}}@inproceedings{Radford2021LearningTV,title={Learning Transferable Visual Models From Natural Language Supervision},author={Alec Radford and Jong Wook Kim and Chris Hallacy and A. Ramesh and Gabriel Goh and Sandhini Agarwal and Girish Sastry and Amanda Askell and Pamela Mishkin and Jack Clark and Gretchen Krueger and Ilya Sutskever},booktitle={ICML},year={2021}}