{"id":21813582,"url":"https://github.com/bobmcdear/open-clip-jax","last_synced_at":"2025-04-13T23:30:45.764Z","repository":{"id":154228361,"uuid":"626581376","full_name":"BobMcDear/open-clip-jax","owner":"BobMcDear","description":"Implementation of OpenAI's CLIP in JAX/Flax","archived":false,"fork":false,"pushed_at":"2024-08-18T18:00:54.000Z","size":1424,"stargazers_count":3,"open_issues_count":0,"forks_count":2,"subscribers_count":1,"default_branch":"main","last_synced_at":"2025-03-27T13:45:21.827Z","etag":null,"topics":["computer-vision","contrastive-learning","deep-learning","flax","jax"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":null,"status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/BobMcDear.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":null,"code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2023-04-11T18:54:13.000Z","updated_at":"2025-01-18T00:30:12.000Z","dependencies_parsed_at":"2024-01-15T15:02:03.251Z","dependency_job_id":"c8304b86-8af5-46ee-beac-6b885e538843","html_url":"https://github.com/BobMcDear/open-clip-jax","commit_stats":null,"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/BobMcDear%2Fopen-clip-jax","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/BobMcDear%2Fopen-clip-jax/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/BobMcDear%2Fopen-clip-jax/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/BobMcDear%2Fopen-clip-jax/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/BobMcDear","download_url":"https://codeload.github.com/BobMcDear/open-clip-jax/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":248796008,"owners_count":21162887,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["computer-vision","contrastive-learning","deep-learning","flax","jax"],"created_at":"2024-11-27T14:30:14.673Z","updated_at":"2025-04-13T23:30:45.708Z","avatar_url":"https://github.com/BobMcDear.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# CLIP in JAX/Flax\n\n• \u003cstrong\u003e[Introduction](#introduction)\u003c/strong\u003e\u003cbr\u003e\n• \u003cstrong\u003e[Installation](#installation)\u003c/strong\u003e\u003cbr\u003e\n• \u003cstrong\u003e[Usage](#usage)\u003c/strong\u003e\u003cbr\u003e\n• \u003cstrong\u003e[Training](#training)\u003c/strong\u003e\u003cbr\u003e\n\u0026nbsp;\u0026nbsp;\u0026nbsp;\u0026nbsp;• \u003cstrong\u003e[Dataset Preparation](#dataset-preparation)\u003c/strong\u003e\u003cbr\u003e\n\u0026nbsp;\u0026nbsp;\u0026nbsp;\u0026nbsp;\u0026nbsp;\u0026nbsp;\u0026nbsp;\u0026nbsp;• \u003cstrong\u003e[CSV](#csv)\u003c/strong\u003e\u003cbr\u003e\n\u0026nbsp;\u0026nbsp;\u0026nbsp;\u0026nbsp;\u0026nbsp;\u0026nbsp;\u0026nbsp;\u0026nbsp;• \u003cstrong\u003e[TFRecord](#tfrecord)\u003c/strong\u003e\u003cbr\u003e\n\u0026nbsp;\u0026nbsp;\u0026nbsp;\u0026nbsp;• \u003cstrong\u003e[Single-Worker Training](#single-worker-training)\u003c/strong\u003e\u003cbr\u003e\n\u0026nbsp;\u0026nbsp;\u0026nbsp;\u0026nbsp;• \u003cstrong\u003e[Multi-Worker Training](#multi-worker-training)\u003c/strong\u003e\u003cbr\u003e\n• \u003cstrong\u003e[Available Models](#available-models)\u003c/strong\u003e\u003cbr\u003e\n• \u003cstrong\u003e[Acknowledgements](#acknowledgements)\u003c/strong\u003e\u003cbr\u003e\n• \u003cstrong\u003e[Citations](#citations)\u003c/strong\u003e\u003cbr\u003e\n\n\n## Introduction\n\n```open_clip_jax``` is an open source JAX/Flax implementation of OpenAI's [CLIP](https://arxiv.org/abs/2103.00020), including image and text towers, pre-trained parameters, training utilities, and more. It is inspired by but not affiliated with [OpenCLIP](https://github.com/mlfoundations/open_clip) and aims to deliver similar functionalities with a JAX backend.\n\n## Installation\n\nThe JAX installation process may differ depending on one's machine, so JAX needs to be [installed manually](https://github.com/google/jax#installation) by the user. Afterwards, ```open_clip_jax``` can be installed through ```pip install git+https://github.com/BobMcDear/open-clip-jax.git```.\n\n## Usage\n\n```CLIPInference``` is a convenience class for conducting inference, which can be called on raw images and texts to compute their similarity scores, as demonstrated below.\n\n```python\nimport jax\nfrom PIL import Image\nfrom open_clip_jax import CLIPInference\n\n\nclip = CLIPInference(\n    'vit-base-patch32',\n    softmax_temp=100.,\n    pretrained='laion2b-s34b-b79k',\n    )\nimage = Image.open('CLIP.png').convert('RGB')\ntext = ['A diagram', 'A dog', 'A cat']\n\n# image and text can be single data points or lists.\nprobs, _ = clip(image, text)\nprint(probs)\n```\n\nUnder the hood, ```CLIPInference``` utilizes ```create_model_with_params``` to create the CLIP model, ```create_image_transforms``` to pre-process the image(s), and ```tokenize``` to tokenize the text(s). A sample usage of these functions, equivalent to the code above, is exhibited in the following snippet. Breaking ```CLIPInference``` into these smaller components can offer greater flexibility.\n\n```python\nfrom typing import Dict\n\nimport jax\nfrom PIL import Image\nfrom jax import Array\nfrom open_clip_jax import create_image_transforms, create_model_with_params, tokenize\n\n\nmodel, vars = create_model_with_params(\n    'vit-base-patch32',\n    pretrained='laion2b-s34b-b79k',\n    )\nimage_transforms = create_image_transforms(\n    train=False,\n    input_format='image',\n    do_batch_transforms=False,\n    )\n\nimage = image_transforms(Image.open('CLIP.png').convert('RGB'))._numpy()\nimage = np.expand_dims(image, axis=0)\ntext = tokenize(['A diagram', 'A dog', 'A cat'])._numpy()\n\ndef calculate_similarity(vars: Dict, image: Array, text: Array) -\u003e Array:\n    # CLIP returns L2-normalized image and text features.\n    image_proj, text_proj = model.apply(vars, image, text)\n    return nn.softmax(100 * image_proj @ text_proj.T)\n\nprobs = jax.jit(calculate_similarity)(vars, image, text)\nprint(probs)\n```\n\n## Training\n\nThis repository also supports training CLIP models from scratch, using either the utilities supplied by ```open_clip_jax.training``` for more fine-grained control or ```main.py``` for a fully-featured training script. The ensuing sections elaborate on training with ```main.py```.\n\n### Dataset Preparation\n\n```main.py``` accepts two data formats, CSV files or TFRecords. The latter should generally be preferred as ```tf.data``` pipelines constructed around TFRecords are quite efficient, especially if the data is stored remotely, but training using CSV files can be more convenient and should not be an issue when dealing with smaller datasets. Dataset preparation instructions for each case are outlined below.\n\n#### CSV\n\nTo prepare a dataset for CSV training, a CSV file needs to be created with one column containing image paths and another holding text captions corresponding to each image. Other columns may be included as well, but they are not read, and the order of the columns is also ignored. For instance, the table below displays how such a file may be structured.\n\n| caption                   |image_path |\n|--------------------------------|------------|\n| Diagram of OpenAI's CLIP model | clip.jpg   |\n| A Siamese cat                  |  cat.jpg    |\n| Dog running on grass           |  dog.jpg    |\n| ...                            | ...        |\n\nTensorFlow integrates seamlessly with Google Cloud Storage (GCS), so the CSV file or images may be stored in a GCS bucket, as can be seen below. However, doing so would slow down data loading since GCS has a high time to first byte (TTFB), and therefore TFRecords would be the appropriate option if storing data in the cloud.\n\n| caption                   |image_path |\n|--------------------------------|------------|\n| Diagram of OpenAI's CLIP model | gs://open_clip_jax/clip.jpg   |\n| A Siamese cat                  |  gs://open_clip_jax/cat.jpg    |\n| Dog running on grass           |  gs://open_clip_jax/dog.jpg    |\n| ...                            | ...        |\n\n\n\n#### TFRecord\n\nTo prepare a dataset for TFRecord training, every image-text pair must be written as a ```tf.train.Example```/Protobuf message (images as JPEG-encoded bytes, text captions as strings) to TFRecord files stored in a single directory (local or in a GCS bucket), with ideally 100+ MB of data, or 10,000 samples, per file. [img2dataset](https://github.com/rom1504/img2dataset) can automatically convert image URLs to such datasets by setting the output format to TFRecord via ```--output_format tfrecord``` and supports many popular datasets, e.g., COCO, LAION-400M, etc.\n\n### Single-Worker Training\n\nIn single-worker settings, assuming JAX and ```open_clip_jax``` have been installed, ```main.py``` simply needs to be downloaded and executed to begin training. If using cloud computing, the remote server should be logged into first, the data optionally transferred to it, and finally the following commands can be run, ideally in a [```tmux```](https://github.com/tmux/tmux/wiki) or [```screen```](https://www.gnu.org/software/screen/) session.\n\n```bash\nwget https://raw.githubusercontent.com/BobMcDear/open-clip-jax/main/main.py -q\npython3 open-clip-jax/main.py \\\n    --train-path train.csv \\\n    --valid-path valid.csv \\\n    --image-key image_path \\\n    --text-key caption \\\n    --global-batch-size 128 \\\n    --model-name vit-base-patch32 \\\n    --learning-rate 1e-3 \\\n    --n-epochs 30\n```\n\n### Multi-Worker Training\n\nMulti-worker training has been tested only for TPUs, although the process should remain largely identical for GPU clusters. To train on a pod slice, the same commands are sent to every worker in parallel to install the necessary packages and start training. The data must also be in a GCS bucket in the same zone as the VM. An end-to-end minimal example, with TPU creation and deletion commands included, can be seen below.\n\n```bash\nNAME=open_clip_jax\nZONE=us-central1-a\nTYPE=v3-32\nVERSION=v2-alpha\n\n# Create pod slice\ngcloud compute tpus tpu-vm create $NAME \\\n    --zone=$ZONE \\\n    --accelerator-type=$TYPE \\\n    --version=$VERSION\n\n# Connect to TPUs and train\ngcloud compute tpus tpu-vm ssh $NAME \\\n    --zone $ZONE \\\n    --worker=all \\\n    --command \"\n        pip install -U pip \u0026\u0026\n        pip install -U jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html \u0026\u0026\n        pip install git+https://github.com/BobMcDear/open-clip-jax.git \u0026\u0026\n        wget https://raw.githubusercontent.com/BobMcDear/open-clip-jax/main/main.py -q \u0026\u0026\n        python3 main.py \\\n            --train-path gs://open_clip_jax/train_tfrecords/ \\\n            --valid-path gs://open_clip_jax/valid_tfrecords/ \\\n            --image-key jpg \\\n            --text-key caption \\\n            --global-batch-size 1024 \\\n            --model-name vit-base-patch32 \\\n            --learning-rate 3e-3 \\\n            --n-epochs 30 \\\n            --checkpoint-dir gs://open_clip_jax/checkpoints/\n        \"\n\n# Delete VM\ngcloud compute tpus tpu-vm delete $NAME \\\n    --zone $ZONE\n```\n\nAn important caveat that should be borne in mind is that epoch boundaries become blurry if the number of batches assigned to each worker varies. For example, suppose a dataset consists of two TFRecord files, one containing 32 samples and the other 64, and training is being performed on two workers (each receiving one TFRecord file) with a per-worker batch size of 16: In this scenario, the number of steps per epoch is calculated to be 3 = (64 + 32) / 32 (the 32 in the denominator is the global batch size), so 3 batches, or 48 = 16 * 3 samples, are taken from each file in every epoch. Consequently, during the initial epoch, half the samples from the first file are iterated over twice, whereas there are samples in the second file that are not seen at all. In the subsequent epoch, however, the remaining samples from the second file will be the first to be grabbed, and the process restarts once the entire file has been consumed. Such extreme discrepancies should be rare as the number of files increases and are unlikely to pose a problem in practice, but one should be cognizant of them nevertheless.\n\n## Available Models\n\nThere are three functions related to listing available models and pre-trained parameters:\n\n* ```list_models```: Returns the name of every model, but some, such as ViT-Small, do not have associated pre-trained parameters.\n* ```list_pretrained```: Returns tuples of (name of model, name of pre-trained parameters). A model may have several groups of pre-trained parameters, so\nthere may be multiple entries with identical model names but different pre-trained parameters.\n* ```list_pretrained_by_model```: Returns a particular model's pre-trained parameters.\n\n```python\n\u003e\u003e\u003e import open_clip_jax\n\u003e\u003e\u003e open_clip_jax.list_models()\n('convnext-base-w',\n 'convnext-base',\n 'convnext-large-d',\n 'vit-base-patch16',\n 'vit-base-patch32',\n 'vit-huge-patch14',\n 'vit-huge-patch16',\n 'vit-large-patch14',\n 'vit-large-patch16',\n 'vit-nano-patch32',\n 'vit-small-patch16',\n 'vit-small-patch32')\n\u003e\u003e\u003e open_clip_jax.list_pretrained()\n(('convnext-base', 'laion400m-s13b-b51k'),\n ('convnext-base-w', 'laion-aesthetic-s13b-b82k'),\n ('convnext-base-w', 'laion-aesthetic-s13b-b82k-320'),\n ('convnext-base-w', 'laion-aesthetic-s13b-b82k-augreg-320'),\n ('convnext-base-w', 'laion2b-s13b-b82k'),\n ('convnext-base-w', 'laion2b-s13b-b82k-augreg'),\n ('convnext-large-d', 'laion2b-s26b-b102k-augreg'),\n ('convnext-large-d', 'laion2b-s29b-b131k-ft-320'),\n ('convnext-large-d', 'laion2b-s29b-b131k-ft-soup-320'),\n ('vit-base-patch32', 'laion400m-e31'),\n ('vit-base-patch32', 'laion400m-e32'),\n ('vit-base-patch32', 'laion2b-e16'),\n ('vit-base-patch32', 'laion2b-s34b-b79k'),\n ('vit-base-patch16', 'laion400m-e31'),\n ('vit-base-patch16', 'laion400m-e32'),\n ('vit-base-patch16', 'laion2b-s34b-b88k'),\n ('vit-large-patch14', 'laion400m-e31'),\n ('vit-large-patch14', 'laion400m-e32'),\n ('vit-large-patch14', 'laion2b-s32b-b82k'),\n ('vit-huge-patch14', 'laion2b-s32b-b79k'))\n\u003e\u003e\u003e open_clip_jax.list_pretrained_by_model('vit-base-patch32')\n('laion400m-e31', 'laion400m-e32', 'laion2b-e16', 'laion2b-s34b-b79k')\n```\n\nThe pre-trained parameters have been ported from OpenCLIP, and more information regarding them, such as their training recipes or zero-shot performance, can be found in the OpenCLIP repository or as model cards on [Hugging Face Hub](https://huggingface.co/models?library=open_clip).\n\n\n## Acknowledgements\n\nThanks to Google's [TPU Research Cloud (TRC) program](https://sites.research.google/trc/about/) for providing hardware used to accelerate the development of this project.\n\n## Citations\n\n```bibtex\n@software{ilharco_gabriel_2021_5143773,\n  author       = {Ilharco, Gabriel and\n                  Wortsman, Mitchell and\n                  Wightman, Ross and\n                  Gordon, Cade and\n                  Carlini, Nicholas and\n                  Taori, Rohan and\n                  Dave, Achal and\n                  Shankar, Vaishaal and\n                  Namkoong, Hongseok and\n                  Miller, John and\n                  Hajishirzi, Hannaneh and\n                  Farhadi, Ali and\n                  Schmidt, Ludwig},\n  title        = {OpenCLIP},\n  month        = jul,\n  year         = 2021,\n  note         = {If you use this software, please cite it as below.},\n  publisher    = {Zenodo},\n  version      = {0.1},\n  doi          = {10.5281/zenodo.5143773},\n  url          = {https://doi.org/10.5281/zenodo.5143773}\n}\n```\n\n```bibtex\n@inproceedings{Radford2021LearningTV,\n  title={Learning Transferable Visual Models From Natural Language Supervision},\n  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},\n  booktitle={ICML},\n  year={2021}\n}\n```\n\n```bibtex\n@inproceedings{schuhmann2022laionb,\n  title={{LAION}-5B: An open large-scale dataset for training next generation image-text models},\n  author={Christoph Schuhmann and\n          Romain Beaumont and\n          Richard Vencu and\n          Cade W Gordon and\n          Ross Wightman and\n          Mehdi Cherti and\n          Theo Coombes and\n          Aarush Katta and\n          Clayton Mullis and\n          Mitchell Wortsman and\n          Patrick Schramowski and\n          Srivatsa R Kundurthy and\n          Katherine Crowson and\n          Ludwig Schmidt and\n          Robert Kaczmarczyk and\n          Jenia Jitsev},\n  booktitle={Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track},\n  year={2022},\n  url={https://openreview.net/forum?id=M3Y74vmsMcY}\n}\n```","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fbobmcdear%2Fopen-clip-jax","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fbobmcdear%2Fopen-clip-jax","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fbobmcdear%2Fopen-clip-jax/lists"}