Package: open-clip-torch
OpenCLIP[Paper][Citations][Clip Colab][Coca Colab]Welcome to an open source implementation of OpenAI'sCLIP(Contrastive Language-Image Pre-training).Using this codebase, we have trained several models on a variety of data sources and compute budgets, ranging fromsmall-scale experimentsto larger runs including models trained on datasets such asLAION-400M,LAION-2BandDataComp-1B.
Many of our models and their scaling properties are studied in detail in the paperreproducible scaling laws for contrastive language-image learning.
Some of our best models and their zero-shot ImageNet-1k accuracy are shown below, along with the ViT-L model trained by OpenAI.
We provide more details about our full collection of pretrained modelshere, and zero-shot results for 38 datasetshere.ModelTraining dataResolution# of samples seenImageNet zero-shot acc.ConvNext-BaseLAION-2B256px13B71.5%ConvNext-LargeLAION-2B320px29B76.9%ConvNext-XXLargeLAION-2B256px34B79.5%ViT-B/32DataComp-1B256px34B72.8%ViT-B/16DataComp-1B224px13B73.5%ViT-L/14LAION-2B224px32B75.3%ViT-H/14LAION-2B224px32B78.0%ViT-L/14DataComp-1B224px13B79.2%ViT-G/14LAION-2B224px34B80.1%ViT-L/14OpenAI's WIT224px13B75.5%Model cards with additional model specific details can be found on the Hugging Face Hub under the OpenCLIP library tag:https://huggingface.co/models?library=open_clip.If you found this repository useful, please considerciting.
We welcome anyone to 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',pretrained='laion2b_s34b_b79k')tokenizer=open_clip.get_tokenizer('ViT-B-32')image=preprocess(Image.open("CLIP.png")).unsqueeze(0)text=tokenizer(["a diagram","a dog","a cat"])withtorch.no_grad(),torch.cuda.amp.autocast():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.]]See also this[Clip Colab].To compute billions of embeddings efficiently, you can useclip-retrievalwhich has openclip support.Pretrained modelsWe offer a simple model interface to instantiate both pre-trained and untrained models.
To see which pretrained models are available, use the following code snippet.
More details about our pretrained models are availablehere.>>>importopen_clip>>>open_clip.list_pretrained()You can find more about the models we support (e.g. number of parameters, FLOPs) inthis table.NOTE: Many existing checkpoints use the QuickGELU activation from the original OpenAI models. This activation is actually less efficient than 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.Loading modelsModels can be loaded withopen_clip.create_model_and_transforms, as shown in the example below. The model name and correspondingpretrainedkeys are compatible with the outputs ofopen_clip.list_pretrained().Thepretrainedargument also accepts local paths, for example/path/to/my/b32.pt.
You can also load checkpoints from huggingface this way. To do so, download theopen_clip_pytorch_model.binfile (for example,https://huggingface.co/laion/CLIP-ViT-L-14-DataComp.XL-s13B-b90K/tree/main), and usepretrained=/path/to/open_clip_pytorch_model.bin.# pretrained also accepts local pathsmodel,_,preprocess=open_clip.create_model_and_transforms('ViT-B-32',pretrained='laion2b_s34b_b79k')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.DataTo download datasets as webdataset, we recommendimg2dataset.Conceptual CaptionsSeecc3m img2dataset example.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 CLIPInstallWe advise you first create a virtual environment with:python3 -m venv .env
source .env/bin/activate
pip install -U pipYou can then install openclip for training withpip install 'open_clip_torch[training]'.DevelopmentIf you want to make changes to contribute code, you can clone openclip then runmake installin openclip folder (after creating a virtualenv)Install pip PyTorch as perhttps://pytorch.org/get-started/locally/You may runmake install-trainingto install training depsTestingTest can be run withmake install-testthenmake testpython -m pytest -x -s -v tests -k "training"to run a specific testRunning regression tests against a specific git revision or tag:Generate testing datapythontests/util_test.py--modelRN50RN101--save_model_listmodels.txt--git_revision9d31b2ec4df6d8228f370ff20c8267ec6ba39383WARNING: This will invoke git and modify your working tree, but will reset it to the current state after data has been generated!Don't modify your working tree while test data is being generated this way.Run regression testsOPEN_CLIP_TEST_REG_MODELS=models.txtpython-mpytest-x-s-v-mregression_testSample 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 does 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.EpochsFor larger datasets (eg Laion2B), we recommend setting--train-num-samplesto a lower value than the full epoch, for example--train-num-samples 135646078to 1/16 of an epoch in conjunction with--dataset-resampledto do sampling with replacement. This allows having frequent checkpoints to evaluate more often.Patch DropoutRecent researchhas shown that one can dropout half to three-quarters of the visual tokens, leading to up to 2-3x training speeds without loss of accuracy.You can set this on your visual transformer config with the keypatch_dropout.In the paper, they also finetuned without the patch dropout at the end. You can do this with the command-line argument--force-patch-dropout 0.Multiple data sourcesOpenCLIP supports using multiple data sources, by separating different data paths with::.
For instance, to train on CC12M and on LAION, one might use--train-data "/data/cc12m/cc12m-train-{0000..2175}.tar::/data/LAION-400M/{00000..41455}.tar".
Using--dataset-resampledis recommended for these cases.By default, on expectation the amount of times the model will see a sample from each source is proportional to the size of the source.
For instance, when training on one data source with size 400M and one with size 10M, samples from the first source are 40x more likely to be seen in expectation.We also support different weighting of the data sources, by using the--train-data-upsampling-factorsflag.
For instance, using--train-data-upsampling-factors=1::1in the above scenario is equivalent to not using the flag, and--train-data-upsampling-factors=1::2is equivalent to upsampling the second data source twice.
If you want to sample from data sources with the same frequency, the upsampling factors should be inversely proportional to the sizes of the data sources.
For instance, if datasetAhas 1000 samples and datasetBhas 100 samples, you can use--train-data-upsampling-factors=0.001::0.01(or analogously,--train-data-upsampling-factors=1::10).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.ptTraining CoCa:TrainingCoCamodels is enabled through specifying a CoCa config using the--modelparameter of the training script. Currently available configs are "coca_base", "coca_ViT-B-32", and "coca_roberta-ViT-B-32" (which uses RoBERTa as the text encoder). CoCa configs are different from CLIP configs because they have an additional "multimodal_cfg" component which specifies parameters for the multimodal text decoder. Here's an example from the coca_ViT-B-32 config:"multimodal_cfg":{"context_length":76,"vocab_size":49408,"width":512,"heads":8,"layers":12,"latent_dim":512,"attn_pooler_heads":8}Credit tolucidrainsforinitial code,gpuccefor adapting the code to open_clip, andiejMacfor training the models.Generating text with CoCaimportopen_clipimporttorchfromPILimportImagemodel,_,transform=open_clip.create_model_and_transforms(model_name="coca_ViT-L-14",pretrained="mscoco_finetuned_laion2B-s13B-b90k")im=Image.open("cat.jpg").convert("RGB")im=transform(im).unsqueeze(0)withtorch.no_grad(),torch.cuda.amp.autocast():generated=model.generate(im)print(open_clip.decode(generated[0]).split("<end_of_text>")[0].replace("<start_of_text>",""))See also this[Coca Colab]Fine Tuning CoCaTo fine-tune coca on mscoco, first create the dataset, one way is using a csvdataset and perhaps the simplest way to do it is usingCLIP_benchmarkwhich in turn usespycocotools(that can be used also by itself).fromclip_benchmark.datasets.builderimportbuild_datasetimportpandasaspdimportosroot_path="path/to/data/dir"# set this to smth meaningfulds=build_dataset("mscoco_captions",root=root_path,split="train")# this downloads the dataset if it is not there alreadycoco=ds.cocoimgs=coco.loadImgs(coco.getImgIds())future_df={"filepath":[],"title":[]}forimginimgs:caps=coco.imgToAnns[img["id"]]forcapincaps:future_df["filepath"].append(img["file_name"])future_df["title"].append(cap["caption"])pd.DataFrame.from_dict(future_df).to_csv(os.path.join(root_path,"train2014.csv"),index=False,sep="\t")This should create a csv dataset that one can use to fine-tune coca with open_clippython-mtraining.main\--dataset-type"csv"\--train-data"path/to/data/dir/train2014.csv"\--warmup1000\--batch-size128\--lr1e-5\--wd0.1\--epochs1\--workers3\--model"coca_ViT-L-14"\--report-to"wandb"\--coca-contrastive-loss-weight0\--coca-caption-loss-weight1\--log-every-n-steps100This is a general setting, open_clip has very parameters that can be set,python -m training.main --helpshould show them. The only relevant change compared to pre-training are the two arguments--coca-contrastive-loss-weight0--coca-caption-loss-weight1which make the model only train the generative side.Training with pre-trained language models as text encoder:If you wish to use different language models as the text encoder for CLIP you can do so by using one of the Hugging Face model configs insrc/open_clip/model_configsand passing in it's tokenizer as the--modeland--hf-tokenizer-nameparameters respectively. Currently we only support RoBERTa ("test-roberta" config), however adding new models should be trivial. You can also determine how many layers, from the end, to leave unfrozen with the--lock-text-unlocked-layersparameter. Here's an example command to train CLIP with the RoBERTa LM that has it's last 10 layers unfrozen:python-mtraining.main\--train-data="pipe:aws s3 cp s3://s-mas/cc3m/{00000..00329}.tar -"\--train-num-samples3000000\--val-data="pipe:aws s3 cp s3://s-mas/cc3m/{00330..00331}.tar -"\--val-num-samples10000\--dataset-typewebdataset\--batch-size256\--warmup2000\--epochs10\--lr5e-4\--precisionamp\--workers6\--model"roberta-ViT-B-32"\--lock-text\--lock-text-unlocked-layers10\--name"10_unfrozen"\--report-to"tensorboard"\Loss 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.LoggingFor tensorboard logging, run:tensorboard--logdir=logs/tensorboard/--port=7777For wandb logging, we recommend looking at thestepvariable instead ofStep, since the later was not properly set in earlier versions of this codebase.
For older runs with models trained beforehttps://github.com/mlfoundations/open_clip/pull/613, theStepvariable should be ignored.
For newer runs, after that PR, the two variables are the same.Evaluation / Zero-ShotWe recommendhttps://github.com/LAION-AI/CLIP_benchmark#how-to-usefor systematic evaluation on 40 datasets.Evaluating 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_e32Model distillationYou can distill from a pre-trained by using--distill-modeland--distill-pretrainedto specify the model you'd like to distill from.
For instance, to distill from OpenAI ViT-L/14 use--distill-model ViT-L-14 --distill-pretrained openai.Gradient accumulationTo simulate larger batches use--accum-freq k. If per gpu batch size,--batch-size, ism, then the effective batch size will bek * m * num_gpus.When increasing--accum-freqfrom its default of 1, samples/s will remain approximately constant (batch size will double, as will time-per-batch). It is recommended to use other features to reduce batch size such as--grad-checkpointing --local-loss --gather-with-gradbefore increasing--accum-freq.--accum-freqcan be used in addition to these features.Instead of 1 forward pass per example, there are now 2 forward passes per-example. However, the first is done withtorch.no_grad.There is some additional GPU memory required --- the features and data from allmbatches are stored in memory.There are alsomloss computations instead of the usual 1.For more information see Cui et al. (https://arxiv.org/abs/2112.09331) or Pham et al. (https://arxiv.org/abs/2111.10050).Int8 SupportWe have beta support for int8 training and inference.
You can enable int8 training with--use-bnb-linear SwitchBackLinearGlobalor--use-bnb-linear SwitchBackLinearGlobalMemEfficient.
Please see the bitsandbytes library for definitions for these layers.
For CLIP VIT-Huge this should currently correspond to a 10% training speedup with no accuracy loss.
More speedups comin when the attention layer is refactored so that linear layers man be replaced there, too.See the tutorialhttps://github.com/mlfoundations/open_clip/blob/main/tutorials/int8_tutorial.ipynborpaper.Support for remote loading/trainingIt is always possible to resume directly from a remote file, e.g., a file in an s3 bucket. Just set--resume s3://<path-to-checkpoint>.
This will work with any filesystem supported byfsspec.It is also possible to trainopen_clipmodels while continuously backing up to s3. This can help to avoid slow local file systems.Say that your node has a local ssd/scratch, an s3 buckets3://<path-to-bucket>.In that case, set--logs /scratchand--remote-sync s3://<path-to-bucket>. Then, a background process will sync/scratch/<run-name>tos3://<path-to-bucket>/<run-name>. After syncing, the background process will sleep for--remote-sync-frequencyseconds, which defaults to 5 minutes.There is also experimental support for syncing to other remote file systems, not just s3. To do so, specify--remote-sync-protocol fsspec. However, this is currently very slow and not recommended.Also, to optionally avoid saving too many checkpoints locally when using these features, you can use--delete-previous-checkpointwhich deletes the previous checkpoint after saving a new one.Note: if you are using this feature with--resume latest, there are a few warnings. First, use with--save-most-recentis not supported. Second, onlys3is supported. Finally, since the sync happens in the background, it is possible that the most recent checkpoint may not be finished syncing to the remote.Pushing Models to Hugging Face HubThe moduleopen_clip.push_to_hf_hubincludes helpers for pushing models /w weights and config to the HF Hub.The tool can be run from command line, ex:python -m open_clip.push_to_hf_hub --model convnext_large_d_320 --pretrained /train/checkpoints/epoch_12.pt --repo-id laion/CLIP-convnext_large_d_320.laion2B-s29B-b131K-ftAcknowledgmentsWe 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,Romain Beaumont,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{cherti2023reproducible,title={Reproducible scaling laws for contrastive language-image learning},author={Cherti, Mehdi and Beaumont, Romain and Wightman, Ross and Wortsman, Mitchell and Ilharco, Gabriel and Gordon, Cade and Schuhmann, Christoph and Schmidt, Ludwig and Jitsev, Jenia},booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},pages={2818--2829},year={2023}}@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}}@inproceedings{schuhmann2022laionb,title={{LAION}-5B: An open large-scale dataset for training next generation image-text models},author={Christoph Schuhmann andRomain Beaumont andRichard Vencu andCade W Gordon andRoss Wightman andMehdi Cherti andTheo Coombes andAarush Katta andClayton Mullis andMitchell Wortsman andPatrick Schramowski andSrivatsa R Kundurthy andKatherine Crowson andLudwig Schmidt andRobert Kaczmarczyk andJenia Jitsev},booktitle={Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track},year={2022},url={https://openreview.net/forum?id=M3Y74vmsMcY}}