GitHub - BobMcDear/open-clip-jax: Implementation of OpenAI's CLIP in JAX/Flax

import jax
from PIL import Image
from open_clip_jax import CLIPInference


clip = CLIPInference(
    'vit-base-patch32',
    softmax_temp=100.,
    pretrained='laion2b-s34b-b79k',
    )
image = Image.open('CLIP.png').convert('RGB')
text = ['A diagram', 'A dog', 'A cat']

# image and text can be single data points or lists.
probs, _ = clip(image, text)
print(probs)

from typing import Dict

import jax
from PIL import Image
from jax import Array
from open_clip_jax import create_image_transforms, create_model_with_params, tokenize


model, vars = create_model_with_params(
    'vit-base-patch32',
    pretrained='laion2b-s34b-b79k',
    )
image_transforms = create_image_transforms(
    train=False,
    input_format='image',
    do_batch_transforms=False,
    )

image = image_transforms(Image.open('CLIP.png').convert('RGB'))._numpy()
image = np.expand_dims(image, axis=0)
text = tokenize(['A diagram', 'A dog', 'A cat'])._numpy()

def calculate_similarity(vars: Dict, image: Array, text: Array) -> Array:
    # CLIP returns L2-normalized image and text features.
    image_proj, text_proj = model.apply(vars, image, text)
    return nn.softmax(100 * image_proj @ text_proj.T)

probs = jax.jit(calculate_similarity)(vars, image, text)
print(probs)

wget https://raw.githubusercontent.com/BobMcDear/open-clip-jax/main/main.py -q
python3 open-clip-jax/main.py \
    --train-path train.csv \
    --valid-path valid.csv \
    --image-key image_path \
    --text-key caption \
    --global-batch-size 128 \
    --model-name vit-base-patch32 \
    --learning-rate 1e-3 \
    --n-epochs 30

NAME=open_clip_jax
ZONE=us-central1-a
TYPE=v3-32
VERSION=v2-alpha

# Create pod slice
gcloud compute tpus tpu-vm create $NAME \
    --zone=$ZONE \
    --accelerator-type=$TYPE \
    --version=$VERSION

# Connect to TPUs and train
gcloud compute tpus tpu-vm ssh $NAME \
    --zone $ZONE \
    --worker=all \
    --command "
        pip install -U pip &&
        pip install -U jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html &&
        pip install git+https://github.com/BobMcDear/open-clip-jax.git &&
        wget https://raw.githubusercontent.com/BobMcDear/open-clip-jax/main/main.py -q &&
        python3 main.py \
            --train-path gs://open_clip_jax/train_tfrecords/ \
            --valid-path gs://open_clip_jax/valid_tfrecords/ \
            --image-key jpg \
            --text-key caption \
            --global-batch-size 1024 \
            --model-name vit-base-patch32 \
            --learning-rate 3e-3 \
            --n-epochs 30 \
            --checkpoint-dir gs://open_clip_jax/checkpoints/
        "

# Delete VM
gcloud compute tpus tpu-vm delete $NAME \
    --zone $ZONE

>>> import open_clip_jax
>>> open_clip_jax.list_models()
('convnext-base-w',
 'convnext-base',
 'convnext-large-d',
 'vit-base-patch16',
 'vit-base-patch32',
 'vit-huge-patch14',
 'vit-huge-patch16',
 'vit-large-patch14',
 'vit-large-patch16',
 'vit-nano-patch32',
 'vit-small-patch16',
 'vit-small-patch32')
>>> open_clip_jax.list_pretrained()
(('convnext-base', 'laion400m-s13b-b51k'),
 ('convnext-base-w', 'laion-aesthetic-s13b-b82k'),
 ('convnext-base-w', 'laion-aesthetic-s13b-b82k-320'),
 ('convnext-base-w', 'laion-aesthetic-s13b-b82k-augreg-320'),
 ('convnext-base-w', 'laion2b-s13b-b82k'),
 ('convnext-base-w', 'laion2b-s13b-b82k-augreg'),
 ('convnext-large-d', 'laion2b-s26b-b102k-augreg'),
 ('convnext-large-d', 'laion2b-s29b-b131k-ft-320'),
 ('convnext-large-d', 'laion2b-s29b-b131k-ft-soup-320'),
 ('vit-base-patch32', 'laion400m-e31'),
 ('vit-base-patch32', 'laion400m-e32'),
 ('vit-base-patch32', 'laion2b-e16'),
 ('vit-base-patch32', 'laion2b-s34b-b79k'),
 ('vit-base-patch32', 'datacomp-m-s128m-b4k'),
 ('vit-base-patch32', 'datacomp-xl-s13b-b90k'),
 ('vit-base-patch16', 'laion400m-e31'),
 ('vit-base-patch16', 'laion400m-e32'),
 ('vit-base-patch16', 'laion2b-s34b-b88k'),
 ('vit-base-patch16', 'dfn2b'),
 ('vit-large-patch14', 'laion400m-e31'),
 ('vit-large-patch14', 'laion400m-e32'),
 ('vit-large-patch14', 'laion2b-s32b-b82k'),
 ('vit-huge-patch14', 'laion2b-s32b-b79k'))
>>> open_clip_jax.list_pretrained_by_model('vit-base-patch32')
('laion400m-e31', 'laion400m-e32', 'laion2b-e16', 'laion2b-s34b-b79k')

@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{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}
}