gem-hunter

[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

OpenCLIP requires Python 3.9+ and PyTorch 2.6+. Python 3.14 and 3.15 have been validated with PyTorch 2.15 for inference, zero-shot evaluation, and compiled training. Install matching PyTorch and torchvision builds for your Python version and hardware.

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