Skip to content

Contrastive Language-Image Pretraining

License

Notifications You must be signed in to change notification settings

nshepperd/CLIP_JAX

This branch is 1 commit ahead of, 3 commits behind kingoflolz/CLIP_JAX:main.

Folders and files

NameName
Last commit message
Last commit date

Latest commit

a858a7e · Jan 26, 2022

History

20 Commits
Mar 1, 2021
Jan 26, 2022
Jan 30, 2021
Jan 26, 2022
Mar 1, 2021
Jan 5, 2021
Jan 5, 2021
Apr 16, 2021
Mar 1, 2021
Jan 12, 2021
Mar 1, 2021
Feb 28, 2021
Mar 1, 2021

Repository files navigation

CLIP (With Haiku + Jax!)

[Blog] [Paper] [Model Card] [Colab]

CLIP (Contrastive Language-Image Pre-Training) is a neural network trained on a variety of (image, text) pairs. It can be instructed in natural language to predict the most relevant text snippet, given an image, without directly optimizing for the task, similarly to the zero-shot capabilities of GPT-2 and 3. We found CLIP matches the performance of the original ResNet50 on ImageNet “zero-shot” without using any of the original 1.28M labeled examples, overcoming several major challenges in computer vision.

Details

The ViT model and checkpoints have been ported to Haiku, while preserving the same output. See tests/test_consistency.py for details.

No JIT/pmap is performed, but pure inference functions for both the text and image encoders are provided from the the clip_jax.load() function which should be easy to run/parallelize how you wish. See test/tpu_bench.py for an example of using pmap.

Usage Example

import numpy as np
from PIL import Image

import clip_jax

image_fn, text_fn, jax_params, jax_preprocess = clip_jax.load('ViT-B/32', "cpu")

image = np.expand_dims(jax_preprocess(Image.open("CLIP.png")), 0)
text = clip_jax.tokenize(["a diagram", "a dog", "a cat"])

image_embed = image_fn(jax_params, image)
text_embed = text_fn(jax_params, text)

TPU performance

On a TPU v3-8 with Jax tpu-vm alpha (test/tpu_bench.py):

10.1361s to compile model
43.9599s for 16 batches
5963.25 examples/s

TODOs

  • Test on TPUs
  • Easier control over precision and device placement
  • Mixed precision training support
  • Support RN50 model

About

Contrastive Language-Image Pretraining

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages

  • Jupyter Notebook 99.2%
  • Other 0.8%