mlfoundations/open_clip

[Paper] [Citations] [Clip Colab] [Coca Colab]

⚠️ Main branch training stack notice

main now uses the post-refactor training stack by default. Training is organized around TrainingTask wrappers, dict-based batches, FSDP2 support, NaFlex image/audio pipelines, and multiple torch.compile strategies. The scope has grown well beyond the original refactor and now includes several new model families. For the older release-stable training API, pin to the v3 branch or the latest 3.x release on PyPI. Inference usage of pretrained image/text models is still intended to be compatible, but training scripts and downstream integrations should review the changes below before upgrading.

New / experimental model families on main:

• NaFlex CLIP — variable-resolution/aspect image towers (timm naflexvit) with token-budget batching (--use-naflex, naflex_* configs)
• NaFlex CLAP — audio-text contrastive training with variable-duration audio (naflexclap_* configs)
• NaFlex GenLIP / GenLAP — generative image/audio captioning with prefix-LM attention and packed [media ; text] rows (naflexgenlip_*, naflexgenlap_* configs, tiktoken text)
• Modern text tower — text_cfg.text_arch="modern": RoPE, SwiGLU/ReLU², RMSNorm, masked pooling (eos/mean/map), with optional qk-norm, gated attention, register tokens, sandwich norm, value residuals, and zero-init residual (moderntext-* configs)
• Variable-length text — text_cfg.variable_text=true pads captions to the per-batch max instead of a fixed context length (works with the modern tower and HF towers)
• Hugging Face ModernBERT text towers — e.g. gte-modernbert-base-ViT-B-32-256
• MaMMUT — single text decoder used in two passes per the paper: bi-directional without cross-attention for contrastive learning, causal with cross-attention for captioning (trains via the CoCa task/loss). Two config families, mirroring coca/coca2: mammut_* reproduces the original LAION-fork numerics exactly via legacy flags (pool_type="avg_all", use_pad_mask=false) and loads the released LAION openMaMMUT-ViT-L-14 weights (pretrained tag datacomp1b_s12_8b_b180k, or directly via hf-hub: — fork-format configs and state dicts are translated on load); mammut2_* uses the corrected defaults (masked-mean text pooling, pad masking). A modern-arch decoder variant is available via multimodal_cfg.text_arch="modern" (mammut2-moderntext_*)
• CoCa v2 configs — coca2_*: paper-faithful attentional pooling (vision_cfg.attentional_pool="cascade", the paper's default — a separate single-query contrastive pooler, so the caption decoder now cross-attends over all 256 generative pooler tokens instead of 255, fixing #458) and the corrected CLS/pad attention mask (text_cfg.correct_cls_mask=true), per #554. A modern-text variant pairs the modern tower with a modern multimodal decoder (coca2-moderntext_*). Existing coca_* configs and released weights are unchanged
• Text validity masks for generative models — CoCa/MaMMUT forward()/encode_text() accept text_valid ([B, L], True = real token); tokenizers can emit exact masks (tokenizer(texts, output_mask=True)), and caption labels are masked to -100 from them. Fixes the SimpleTokenizer pad-collision class (id 0 is a real token, '!' merges like x!=y emit it mid-caption); absent a mask, behavior falls back to the historical text != pad_id derivation. Text towers keep the HF-style attention_mask kwarg at the tower boundary
• Variable text with gradient accumulation — CoCa/MaMMUT support different caption lengths across microbatches in both training entrypoints. Model forwards retain each microbatch's text length; caption logits and masked labels are padded when combining loss inputs, preserving the mean over all valid target tokens. This applies to the logits-based caption loss; --fused-caption-loss still requires --accum-freq 1.

Breaking changes — training CLI:

• --horovod removed (Horovod support deleted; DDP/FSDP2 only)
• --torchscript and --trace removed (torch.jit is being deprecated upstream)
• Default --precision changed from amp → amp_bf16 (silent behavior change; pass --precision amp explicitly to keep fp16 AMP)
• SigLIP's --loss-dist-impl now defaults to gather, as does standalone SigLipLoss. Pass --loss-dist-impl bidir to keep bidirectional ring exchange; reduce and shift remain available. Gather stores all ranks' text features on each rank.
• --naflex-max-tokens-per-batch now defaults to unset. The local token budget is inferred as --batch-size * max(--naflex-seq-lens); GenLIP/GenLAP also include their caption-token cap in the per-row cost. Pass an explicit token budget to preserve older runs that relied on the previous 16384 default.

New training CLI flags (opt-in):

• --siglip-chunk-size — image rows per SigLIP logits chunk (0 disables). For example, --siglip --siglip-chunk-size 1024 enables chunking in both current and legacy training.
• --fsdp — use FSDP2 (fully_shard) instead of DDP
• --fsdp-no-reshard-after-forward, --fsdp-offload-cpu
• --fsdp-checkpoint {full,sharded} — full gathers to rank-0 as a single .pt; sharded uses DCP per-rank shards (faster, lower memory)
• --torchcompile-strategy {task,model,step} — choose whether torch.compile captures task forward/loss, the underlying model, or the full single-batch train step
• --use-naflex and --naflex-* flags — enable NaFlex variable-aspect image pipelines for compatible timm/OpenCLIP ViT-family models (token-budget batching via --naflex-seq-lens / --naflex-max-tokens-per-batch; the same machinery drives the NaFlex audio and generative models)
• --audio-* and --audio-zeroshot-* flags — enable CLAP audio preprocessing/training and Hugging Face audio zero-shot evaluation
• --length-bucketing, --bucket-pool, --bucket-chunk — reorder the train stream by sample length (caption and/or audio tokens) to tighten per-batch padding; the bucket pool holds raw, undecoded samples
• --text-pad-multiple — round per-batch variable-text length up to a multiple, bounding the distinct sequence lengths torch.compile sees (text-axis analogue of --naflex-pad-multiple)
• --text-attention-mask — emit a per-sample text validity mask (batch key text_valid) from the tokenizer, consumed by CoCa/MaMMUT for attention/pooling and -100 caption-label masking. Default auto-enables for CoCa/MaMMUT (except under --distill) and is rejected for tasks that don't consume it
• --caption-z-loss-weight, --caption-loss-compute-dtype {float32,model}, and --caption-loss-chunk-size — configure the next-token objective shared by CoCa/MaMMUT and GenLIP/GenLAP. Defaults preserve the existing fp32 CE with no z-loss; model preserves the loss-logit dtype and ambient autocast policy while returned loss/component scalars remain fp32

Breaking changes — Python API:

• trace_model removed from the top-level open_clip namespace
• load_openai_model, list_openai_models, and build_model_from_openai_state_dict removed. Original-OpenAI weights are still loadable through the standard create_model_from_pretrained(..., pretrained='openai') path, which now routes through HuggingFace Hub (timm/*_clip.openai) instead of torch.jit.load on openaipublic.azureedge.net archives. Removing the JIT path closes an arbitrary-code-execution surface (JIT archives can ship pickled code).
• Training pipeline wraps model + loss in a TrainingTask subclass (CLIPTask, SigLIPTask, CoCaTask, DistillCLIPTask, CLAPTask). Code that previously called train_one_epoch(model, loss, ...) or evaluate(model, ...) directly should switch to passing a task. Tasks construct their own losses; create_loss("clip", ...) is available for standalone training loops (see below).
• Data pipelines emit dict batches instead of tuples. Image/text loaders use {"image": ..., "text": ...}; CLAP audio loaders use {"audio": ..., "text": ...}. Tuple-style image/text calls through task(images, texts) still work via a backward-compat path, but downstream code that iterates a dataloader directly should read the named keys.
• CoCa's autoregressive label shift moved out of coca_model.py and into CoCaTask. coca_model.forward() no longer performs the [:, :-1] / [:, 1:] slicing — callers that relied on that behavior outside training should handle labels themselves.
• CoCa's exact-mask API adds text_valid after text in encode_text, forward, and forward_intermediates. This shifts the older trailing positional arguments (normalize, image_latent, image_indices, and so on); pass those arguments by keyword. For example, replace model.encode_text(text, False) with model.encode_text(text, normalize=False).
• CLIPTextCfg.eos_id no longer defaults to 2 (that value is only correct for XLM-style vocabs). Configs using pool_type="eos" must set eos_id explicitly, and get_tokenizer now validates eos_id/pad_id against the resolved tokenizer, raising on mismatch instead of pooling/masking silently wrong positions.
• HFTokenizer no longer fabricates pad_token_id=0 when the underlying tokenizer has no pad token (id 0 is a real token in most BPE vocabs); variable-text setups fail fast instead. It also forces padding_side='right', which all OpenCLIP pooling/masking assumes.
• Tokenizer wrappers now share special-token controls: encode(..., add_special_tokens=False) remains body-only by default, while model-facing tokenizer(...) defaults to add_special_tokens=True. decode() / batch_decode() default to skip_special_tokens=False, stop_at_eos=True; pass stop_at_eos=False to inspect tokens after the first EOS. This intentionally changes legacy SimpleTokenizer decode output by hiding post-EOS id-0 fill (!) and makes TikToken decode render its reserved control tokens unless skip_special_tokens=True.
• Model traits replace model-name sniffing. Every model carries model.traits (open_clip.get_model_traits(model)): family, objectives, NaFlex/variable-text contracts. The training entrypoints derive NaFlex data, variable text, the --text-attention-mask default and the grad-accum / distill guards from the built model via apply_model_traits, so hf-hub: and renamed configs no longer need a magic substring; args.genlip / args.genlap / args.naflexclap are gone and the data loaders (get_data, get_wds_dataset, ...) take model_traits. --use-naflex now always sets force_naflex_vision, which the factory treats as a no-op for NaFlex-native and audio models. GenLIP configs get NaFlex transforms from a plain create_model_and_transforms call. create_task() likewise selects the task from the built model. MaMMUT decoders now honor multimodal_cfg.variable_text (previously dropped), so mammut2-moderntext_* configs train with per-batch padded text.
• MaxPooler (hf_pooler_type="max_pooler") mask polarity fixed — it previously max-pooled over the padding positions instead of the valid ones.
• CoCa.__init__ no longer takes a pad_id argument — model.pad_id is derived from the text tower (the id it actually masks with: text_cfg.pad_id for native towers, the transformers config pad for HF towers). MaMMUT follows the same pattern. This fixes coca_roberta-*, which previously masked with roberta's pad (1) in the tower while the loss ignored 0 — its config now declares pad_id: 1 and the caption loss no longer trains on padding.
• CoCaTask builds caption labels masked to -100; CoCaLoss's cross-entropy uses ignore_index=-100. Its pad_id arg is retained for standalone callers passing raw labels (default 0 preserves the old value-based behavior; the task path passes None). Validation generative-loss metrics are likewise mask/pad-aware and will report different (correct) values for nonzero-pad models.
• Standalone loss factory: create_loss(loss_type, *, ...) takes an explicit loss type ("clip", "distill_clip", "siglip", "coca", or "genlip") and flat keyword options, with no model, traits, or CLI namespace required. For example, create_loss("siglip", chunk_size=1024) or create_loss("coca", caption_loss_weight=2.0, clip_loss_weight=1.0, pad_id=1). CoCa/MaMMUT callers must explicitly pass a pad ID for raw labels or pad_id=None for labels already masked to -100; caption logits and labels must already be aligned for next-token prediction. Non-default options unsupported by the selected loss raise. The legacy trainer preserves its create_loss(args, model=...) call through open_clip_train.loss.create_loss_from_args, which resolves model traits, the actual pad ID, and compile-aware label caching. There is no model-name fallback.
• --lock-image now freezes vision attentional poolers (they were previously left trainable when locking CoCa-style towers, and were ungrouped under --image-layer-decay).
• MultimodalTransformer (the CoCa text decoder) now initializes its parameters — previously its init_parameters was never called and text_projection was uninitialized memory on fresh builds. New training runs won't bit-match runs started under older code; checkpoint loading is unaffected. generate() for CoCa/MaMMUT now derives pad/bos/eos ids from the model/tower/HF config instead of hardcoding CLIP token ids (fixes generation for HF-tower CoCa), and shares one implementation in open_clip.generation.
• get_tokenizer validation is stricter for generative configs (CoCa/MaMMUT/GenLIP): an explicit nonzero pad_id with a tokenizer that reserves no pad token raises, and an unset pad_id against a reserved nonzero tokenizer pad warns (the pad-value fallback would train the caption loss on padding).
• WebDataset pipelines now assemble as tokenize -> [length bucketing] -> decode -> transform so the bucket pool holds raw bytes instead of decoded images/waveforms (10-50x less dataloader-worker memory). The old decode-first assembly is preserved in open_clip_train.legacy_data (no bucketing/NaFlex) and used by legacy_main.

Dependency bump: Minimum torch>=2.6 (was >=2.0). This is the version where torch.load(weights_only=True) became the default — all checkpoint loads in this repo now pass weights_only=True explicitly with no weights_only=False fallback. If you're resuming training with a custom optimizer that pickles non-allowlisted Python types, register them via torch.serialization.add_safe_globals([...]) before calling load_checkpoint.

Checkpoint compatibility: Existing pretrained .pt checkpoints load without changes. Training checkpoints saved on main include a state_dict key that's compatible with prior versions; EMA, optimizer state, and optional training counters are also preserved. 0-D vs 1-D scalar reshape from the FSDP path is reconciled on load, so you can resume a DDP-trained checkpoint under FSDP2 and vice versa.

Legacy training entry point: python -m open_clip_train.legacy_main remains available for older image/text training scripts that need the pre-task loop, and pairs with the frozen decode-first data pipelines in open_clip_train.legacy_data. It does not support the full task-era feature set (for example FSDP2, EMA, CLAP audio training, NaFlex, length bucketing, or the task/step torch.compile integration), and should be treated as a compatibility shim rather than the path for new training work.

Welcome to an open source implementation of OpenAI's CLIP (Contrastive Language-Image Pre-training).

Using this codebase, we have trained several models on a variety of data sources and compute budgets, ranging from small-scale experiments to larger runs including models trained on datasets such as LAION-400M, LAION-2B and DataComp-1B. Many of our models and their scaling properties are studied in detail in the paper reproducible scaling laws for contrastive language-image learning. Some of the best models we've trained and their zero-shot ImageNet-1k accuracy are shown below, along with the ViT-L model trained by OpenAI and other state-of-the-art open source alternatives (all can be loaded via OpenCLIP). We provide more details about our full collection of pretrained models here, and zero-shot results for 38 datasets here. Experimental helper scripts for multilabel zero-shot evaluation (scripts/multilabel_zeroshot.py) and linear probes (scripts/linear_probe.py) are also available.

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 consider citing. We welcome anyone to submit an issue or send an email if you have any other requests or suggestions.

Note that portions of src/open_clip/ modelling and tokenizer code are adaptations of OpenAI's official repository.

Image Credit: https://github.com/openai/CLIP

pip install open_clip_torch

import torch
from PIL import Image
import open_clip

model, _, preprocess = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion2b_s34b_b79k')
model.eval()  # model in train mode by default, impacts some models with BatchNorm or stochastic depth active
tokenizer = open_clip.get_tokenizer('ViT-B-32')

image = preprocess(Image.open("docs/CLIP.png")).unsqueeze(0)
text = tokenizer(["a diagram", "a dog", "a cat"])

with torch.no_grad(), torch.autocast("cuda"):
    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.]]

If model uses timm image encoders (convnext, siglip, eva, etc) ensure the latest timm is installed. Upgrade timm if you see 'Unknown model' errors for the image encoder.

If model uses transformers tokenizers, ensure transformers is installed.

See also this [Clip Colab].

To compute billions of embeddings efficiently, you can use clip-retrieval which has openclip support.

We 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 available here.

>>> import open_clip
>>> open_clip.list_pretrained()

You can find more about the models we support (e.g. number of parameters, FLOPs) in this 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 -quickgelu postfix for the OpenCLIP pretrained weights. All OpenAI pretrained weights will always default to QuickGELU. One can also use the non -quickgelu model 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.

Models can be loaded with open_clip.create_model_and_transforms, as shown in the example below. The model name and corresponding pretrained keys are compatible with the outputs of open_clip.list_pretrained().

The pretrained argument also accepts local paths, for example /path/to/my/b32.pt. You can also load checkpoints from huggingface this way. To do so, download the open_clip_pytorch_model.bin file (for example, https://huggingface.co/laion/CLIP-ViT-L-14-DataComp.XL-s13B-b90K/tree/main), and use pretrained=/path/to/open_clip_pytorch_model.bin.

# pretrained also accepts local paths
model, _, preprocess = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion2b_s34b_b79k') 

This repository is focused on training CLIP models. To fine-tune a trained zero-shot model on a downstream classification task such as ImageNet, please see our other repository: WiSE-FT. The WiSE-FT repository contains code for our paper on Robust Fine-tuning of Zero-shot Models, in which we introduce a technique for fine-tuning zero-shot models while preserving robustness under distribution shift.

To download datasets as webdataset, we recommend img2dataset.

See cc3m img2dataset example.

In addition to specifying the training data via CSV files as mentioned above, our codebase also supports webdataset, which is recommended for larger scale datasets. The expected format is a series of .tar files. Each of these .tar files 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.tar could contain files such as abc.jpg and abc.txt. You can learn more about webdataset at https://github.com/webdataset/webdataset. We use .tar files with 1,000 data points each, which we create using tarp.

You can download the YFCC dataset from Multimedia Commons. Similar to OpenAI, we used a subset of YFCC to reach the aforementioned accuracy numbers. The indices of images in this subset are in OpenAI's CLIP repository.

We advise you first create a virtual environment with:

python3 -m venv .env
source .env/bin/activate
pip install -U pip

You can then install openclip for training with pip install 'open_clip_torch[training]'.

If you want to make changes to contribute code, you can clone openclip then run make install in openclip folder (after creating a virtualenv)

Install pip PyTorch as per https://pytorch.org/get-started/locally/

You may run make install-training to install training deps

Test can be run with make install-test then make test

python -m pytest -x -s -v tests -k "training" to run a specific test

Running regression tests against a specific git revision or tag:

1. Generate testing data

   python tests/util_test.py --model RN50 RN101 --save_model_list models.txt --git_revision 9d31b2ec4df6d8228f370ff20c8267ec6ba39383

   WARNING: 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.

2. Run regression tests

   OPEN_CLIP_TEST_REG_MODELS=models.txt python -m pytest -x -s -v -m regression_test

python -m open_clip_train.main \
    --save-frequency 1 \
    --zeroshot-frequency 1 \
    --report-to tensorboard \
    --train-data="/path/to/train_data.csv"  \
    --val-data="/path/to/validation_data.csv"  \
    --csv-img-key filepath \
    --csv-caption-key title \
    --imagenet-val=/path/to/imagenet/root/val/ \
    --warmup 10000 \
    --batch-size=128 \
    --lr=1e-3 \
    --wd=0.1 \
    --epochs=30 \
    --workers=8 \
    --model RN50

Note: imagenet-val is the path to the validation set 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 the val folder should contain subfolders. If it does not, please use this script.

Training supports torch.compile via --torchcompile. Use --torchcompile-strategy task for the common case: it compiles the task train/eval forward callables while keeping the TrainingTask object itself unwrapped for checkpointing, batch preparation, eval, and zero-shot helpers. Use --torchcompile-strategy model when you want the lower-risk pre-DDP behavior of compiling only the underlying trainable module, especially for fp16 AMP runs that need a GradScaler. Use --torchcompile-strategy step for the most aggressive single-batch path: it compiles forward, loss, backward, optional grad clipping, optimizer step, and logit-scale clamp, but requires --accum-freq 1 and a precision mode without GradScaler such as the default amp_bf16 or fp32. --torchcompile-backend and --torchcompile-mode are passed through to torch.compile.

This 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 be O(n^2). Instead, complexity may become effectively linear if the flags --gather-with-grad and --local-loss are used. This alteration results in one-to-one numerical results as the naïve method.

For larger datasets (eg Laion2B), we recommend setting --train-num-samples to a lower value than the full epoch, for example --train-num-samples 135646078 to 1/16 of an epoch in conjunction with --dataset-resampled to do sampling with replacement. This allows having frequent checkpoints to evaluate more often.

Recent research has 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 key patch_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.

OpenCLIP 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-resampled is 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-factors flag. For instance, using --train-data-upsampling-factors=1::1 in the above scenario is equivalent to not using the flag, and --train-data-upsampling-factors=1::2 is 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 dataset A has 1000 samples and dataset B has 100 samples, you can use --train-data-upsampling-factors=0.001::0.01 (or analogously, --train-data-upsampling-factors=1::10).

We make use of torchrun to launch distributed jobs. The following launches a a job on a node of 4 GPUs:

cd open_clip/src
torchrun --nproc_per_node 4 -m open_clip_train.main \
    --train-data '/data/cc12m/cc12m-train-{0000..2175}.tar' \
    --train-num-samples 10968539 \
    --dataset-type webdataset \
    --batch-size 320 \
    --precision amp \
    --workers 4 \
    --imagenet-val /data/imagenet/validation/

The same script above works, so long as users include information about the number of nodes and host node.

cd open_clip/src
torchrun --nproc_per_node=4 \
    --rdzv_endpoint=$HOSTE_NODE_ADDR \
    -m open_clip_train.main \
    --train-data '/data/cc12m/cc12m-train-{0000..2175}.tar' \
    --train-num-samples 10968539 \
    --dataset-type webdataset \
    --batch-size 320 \
    --precision amp \
    --workers 4 \
    --imagenet-val /data/imagenet/validation/

This 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_NAME

eval "$(/path/to/conda/bin/conda shell.bash hook)" # init conda
conda activate open_clip
export CUDA_VISIBLE_DEVICES=0,1,2,3
export MASTER_PORT=12802

master_addr=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
export MASTER_ADDR=$master_addr

cd /shared/open_clip
export PYTHONPATH="$PYTHONPATH:$PWD/src"
srun --cpu_bind=v --accel-bind=gn python -u src/open_clip_train/main.py \
    --save-frequency 1 \
    --report-to tensorboard \
    --train-data="/data/LAION-400M/{00000..41455}.tar" \
    --warmup 2000 \
    --batch-size=256 \
    --epochs=32 \
    --workers=8 \
    --model ViT-B-32 \
    --name "ViT-B-32-Vanilla" \
    --seed 0 \
    --local-loss \
    --gather-with-grad

python -m open_clip_train.main \
    --train-data="/path/to/train_data.csv" \
    --val-data="/path/to/validation_data.csv"  \
    --resume /path/to/checkpoints/epoch_K.pt

Training CoCa models is enabled through specifying a CoCa config using the --model parameter 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 to lucidrains for initial code, gpucce for adapting the code to open_clip, and iejMac for training the models.

import open_clip
import torch
from PIL import Image

model, _, 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)

with torch.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]

To fine-tune coca on mscoco, first create the dataset, one way is using a csvdataset and perhaps the simplest way to do it is using CLIP_benchmark which in turn uses pycocotools (that can be used also by itself).

from clip_benchmark.datasets.builder import build_dataset
import pandas as pd
import os

root_path = "path/to/data/dir" # set this to smth meaningful
ds = build_dataset("mscoco_captions", root=root_path, split="train", task="captioning") # this downloads the dataset if it is not there already
coco = ds.coco
imgs = coco.loadImgs(coco.getImgIds())
future_df = {"filepath":[], "title":[]}
for img in imgs:
    caps = coco.imgToAnns[img["id"]]
    for cap in caps:
        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_clip

python -m open_clip_train.main \
    --dataset-type "csv" \
    --train-data "path/to/data/dir/train2014.csv" \
    --warmup 1000 \
    --batch-size 128 \
    --lr 1e-5 \
    --wd 0.1 \
    --epochs 1 \
    --workers 3 \
    --model "coca_ViT-L-14" \
    --report-to "wandb" \
    --coca-contrastive-loss-weight 0 \
    --coca-caption-loss-weight 1 \
    --log-every-n-steps 100

This is a general setting, open_clip has very parameters that can be set, python -m open_clip_train.main --help should show them. The only relevant change compared to pre-training are the two arguments

--coca-contrastive-loss-weight 0
--coca-caption-loss-weight 1

which make the model only train the generative side.

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 in src/open_clip/model_configs and passing in it's tokenizer as the --model and --hf-tokenizer-name parameters 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-layers parameter. Here's an example command to train CLIP with the RoBERTa LM that has it's last 10 layers unfrozen:

python -m open_clip_train.main \
         --train-data="pipe:aws s3 cp s3://s-mas/cc3m/{00000..00329}.tar -" \
         --train-num-samples 3000000 \
         --val-data="pipe:aws s3 cp s3://s-mas/cc3m/{00330..00331}.tar -" \
         --val-num-samples 10000 \
         --dataset-type webdataset \
         --batch-size 256 \
         --warmup 2000 \
         --epochs 10 \
         --lr 5e-4 \
         --precision amp \
         --workers 6 \
         --model "roberta-ViT-B-32" \
         --lock-text \
         --lock-text-unlocked-layers 10 \
         --name "10_unfrozen" \
         --report-to "tensorboard" \

When 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 of lr=5e-4 and epochs=32.

Note that to use another model, like ViT-B/32 or RN50x4 or RN50x16 or ViT-B/16, specify with --model RN50x4.

For tensorboard logging, run:

tensorboard --logdir=logs/tensorboard/ --port=7777

For wandb logging, we recommend looking at the step variable instead of Step, since the later was not properly set in earlier versions of this codebase. For older runs with models trained before #613, the Step variable should be ignored. For newer runs, after that PR, the two variables are the same.

We recommend https://github.com/LAION-AI/CLIP_benchmark#how-to-use for systematic evaluation on 40 datasets.

python -m open_clip_train.main \
    --val-data="/path/to/validation_data.csv"  \
    --model RN101 \
    --pretrained /path/to/checkpoints/epoch_K.pt

python -m open_clip_train.main \
    --imagenet-val /path/to/imagenet/validation \
    --model ViT-B-32-quickgelu \
    --pretrained laion400m_e32

Audio zero-shot evaluation uses Hugging Face audio classification datasets. Install the optional audio dependencies before running these examples.

pip install 'datasets[audio]' torchaudio torchlibrosa

ESC-50:

python -m open_clip_train.main \
    --model CLAP-HTSAT-tiny-Roberta-base-fused \
    --pretrained laion \
    --audio-zeroshot-dataset ashraq/esc50 \
    --audio-zeroshot-split train \
    --batch-size 64 \
    --workers 4 \
    --device cuda \
    --zeroshot-frequency 1

UrbanSound8K:

python -m open_clip_train.main \
    --model CLAP-HTSAT-tiny-Roberta-base-fused \
    --pretrained laion \
    --audio-zeroshot-dataset mteb/urbansound8K \
    --audio-zeroshot-split train \
    --audio-zeroshot-target-key classID \
    --audio-zeroshot-class-key class \
    --batch-size 64 \
    --workers 4 \
    --device cuda \
    --zeroshot-frequency 1

You can distill from a pre-trained by using --distill-model and --distill-pretrained to 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.

To simulate larger batches use --accum-freq k. If per gpu batch size, --batch-size, is m, then the effective batch size will be k * m * num_gpus.

When increasing --accum-freq from 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-grad before increasing --accum-freq. --accum-freq can 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 with torch.no_grad.

There is some additional GPU memory required --- the features and data from all m batches are stored in memory.

There are also m loss 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).

We have beta support for int8 inference via bitsandbytes' Linear8bitLt (the LLM.int8() layer). Only the MLP linear layers (c_fc/c_proj) are quantized by default; attention's combined QKV projection isn't an nn.Linear submodule and isn't covered yet (see the TODO in src/open_clip/utils.py).

from functools import partial
import bitsandbytes as bnb
import open_clip

model, _, preprocess = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion2b_s34b_b79k')
model = model.half()

int8_linear_layer = partial(bnb.nn.Linear8bitLt, has_fp16_weights=False)
int8_model = open_clip.utils.replace_linear(model, int8_linear_layer, include_modules=['c_fc', 'c_proj']).cuda()
open_clip.utils.convert_int8_model_to_inference_mode(int8_model)  # stamps the dtype bookkeeping the model needs post-quantization

Accuracy / speed, measured, not estimated: on ViT-B-32-quickgelu (openai weights), zero-shot classification on CIFAR-10 was 88.76% (fp16) vs. 88.83% (int8) — a +0.07 point difference, i.e. no meaningful accuracy cost for quantizing just the MLP layers post-hoc (no fine-tuning). This is a single-template, single-dataset, single-model measurement, not a comprehensive benchmark, but it's real and reproducible (see the tutorial notebook). Speed was slightly worse for int8 (53.9ms vs. 56.9ms for a batch of 128, about 5.6% slower) — at this model scale, the extra kernel launches int8 quantization/dequantization requires per forward pass outweigh any raw int8 tensor-core throughput advantage, since Linear8bitLt was designed around much larger (billion-parameter) models than CLIP's transformer blocks. The practical benefit here is memory (roughly a 2x reduction on the quantized layers' weight tensors), not speed or accuracy.

Saving and loading a quantized model: replace_linear swaps nn.Linear submodules for bitsandbytes layers in place, so a state_dict saved from the quantized model only loads back correctly into a model that already has the same layers swapped in — build the model, replace_linear, then torch.save(int8_model.state_dict(), path); to load, rebuild the architecture, replace_linear + convert_int8_model_to_inference_mode again, then load_state_dict. The image transform and tokenizer aren't affected by quantization at all — they only depend on model_name, so just recreate them normally with create_model_and_transforms/get_tokenizer. See the tutorial notebook for a runnable example.

Older SwitchBackLinear/triton path (training): the previous --use-bnb-linear SwitchBackLinearGlobal / SwitchBackLinearGlobalMemEfficient training path, and the triton-based classes in bitsandbytes.nn.triton_based_modules the earlier version of this tutorial used, are effectively unusable on a fresh install today: triton==2.0.0.post1 (which those kernels were written against) has been removed from PyPI entirely, every installable triton version reorganized triton.language.libdevice away, and bitsandbytes>=0.50 dropped triton_based_modules outright. If you hit AttributeError: module 'triton.language' has no attribute 'libdevice', that's why — it's not fixable by pinning a version anymore. Use Linear8bitLt (above) for inference. There's currently no equivalent replacement documented for the training-time SwitchBack speedup path.

See the tutorial https://github.com/mlfoundations/open_clip/blob/main/tutorials/int8_tutorial.ipynb, and the original SwitchBack paper for the training-time approach (0.1pt accuracy cost vs. bf16 on CLIP ViT-Huge, per the paper — a very different regime from the post-hoc inference-only quantization measured above).

It 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 by fsspec.

It is also possible to train open_clip models 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 bucket s3://<path-to-bucket>.

In that case, set --logs /scratch and --remote-sync s3://<path-to-bucket>. Then, a background process will sync /scratch/<run-name> to s3://<path-to-bucket>/<run-name>. After syncing, the background process will sleep for --remote-sync-frequency seconds, 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-checkpoint which 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-recent is not supported. Second, only s3 is 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.

The module open_clip.push_to_hf_hub includes 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-ft

We 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).

Current development of this repository is led by Ross Wightman, Romain Beaumont, Cade Gordon, and Vaishaal 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 Schmidt

Special thanks to Jong Wook Kim and Alec Radford for help with reproducing CLIP!

If you found this repository useful, please consider citing:

@software{ilharco_gabriel_2021_5143773,
  author       = {Ilharco, Gabriel and
                  Wortsman, Mitchell and
                  Wightman, Ross and
                  Gordon, Cade and
                  Carlini, Nicholas and
                  Taori, Rohan and
                  Dave, Achal and
                  Shankar, Vaishaal and
                  Namkoong, Hongseok and
                  Miller, John and
                  Hajishirzi, Hannaneh and
                  Farhadi, Ali and
                  Schmidt, 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 and
          Romain Beaumont and
          Richard Vencu and
          Cade W Gordon and
          Ross Wightman and
          Mehdi Cherti and
          Theo Coombes and
          Aarush Katta and
          Clayton Mullis and
          Mitchell Wortsman and
          Patrick Schramowski and
          Srivatsa R Kundurthy and
          Katherine Crowson and
          Ludwig Schmidt and
          Robert Kaczmarczyk and
          Jenia Jitsev},
  booktitle={Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track},
  year={2022},
  url={https://openreview.net/forum?id=M3Y74vmsMcY}
}