x-clip

To support multiview contrastive learning, set multiview_loss_weight during CLIP initialization. When calling the model, pass augmented text and/or images using aug_text and aug_image. These arguments can accept single tensors or tuples of tensors for multiple augmentations.

• multiview_loss_weight: Weight for the multiview contrastive loss.
• aug_text: Augmented text (e.g., backtranslation). Shape: (batch, seq_len) or tuple of such tensors.
• aug_image: Augmented images. Shape: (batch, C, H, W) or tuple of such tensors.

import torch
from x_clip import CLIP, TextTransformer
# ... (setup encoders) ...

clip = CLIP(
    image_encoder = image_encoder,
    text_encoder = text_encoder,
    dim_image = 512,
    dim_text = 512,
    dim_latent = 512,
    extra_latent_projection = True,
    multiview_loss_weight = 0.1
)

text = torch.randint(0, 10000, (4, 256))
images = torch.randn(4, 3, 256, 256)
aug_text = torch.randint(0, 10000, (4, 256))
aug_images = torch.randn(4, 3, 256, 256)

loss = clip(
    text,
    images,
    aug_text = aug_text,
    aug_image = aug_images,
    return_loss = True,
    freeze_image_encoder = True
)

loss.backward()

Initialize the CLIP class with various hyperparameters to define the dimensions and architecture of the text and image encoders. You can then pass text tokens and images to the instance to compute the contrastive loss.

Key initialization parameters:

• dim_text, dim_image, dim_latent: Dimensions for text, image, and latent embeddings.
• num_text_tokens: Vocabulary size.
• visual_patch_dropout: Patch dropout probability (e.g., 0.5 for FLIP).
• use_all_token_embeds: Enables fine-grained contrastive learning (FILIP).
• decoupled_contrastive_learning: Enables DCL objective (CLOOB + DCL).
• extra_latent_projection: Enables separate projections for text-to-image vs image-to-text (CLOOB).
• use_visual_ssl: Enables self-supervised learning on images.
• use_mlm: Enables masked language learning on text (DeCLIP).

When calling the CLIP instance, use return_loss=True to obtain the loss tensor for backpropagation.

import torch
from x_clip import CLIP

clip = CLIP(
    dim_text = 512,
    dim_image = 512,
    dim_latent = 512,
    num_text_tokens = 10000,
    text_enc_depth = 6,
    text_seq_len = 256,
    text_heads = 8,
    visual_enc_depth = 6,
    visual_image_size = 256,
    visual_patch_size = 32,
    visual_heads = 8,
    visual_patch_dropout = 0.5,
    use_all_token_embeds = False,
    decoupled_contrastive_learning = True,
    extra_latent_projection = True,
    use_visual_ssl = True,
    use_mlm = False,
    text_ssl_loss_weight = 0.05,
    image_ssl_loss_weight = 0.05
)

# mock data
text = torch.randint(0, 10000, (4, 256))
images = torch.randn(4, 3, 256, 256)

# train
loss = clip(
    text,
    images,
    freeze_image_encoder = False,
    return_loss = True
)

loss.backward()

You can provide a custom image encoder to the CLIP class. The encoder must return embeddings in the shape batch x seq x dim. Ensure that dim_image in the CLIP constructor matches the dimension of the returned embeddings.

import torch
from x_clip import CLIP
from vit_pytorch import ViT
from vit_pytorch.extractor import Extractor

# Requires: pip install vit_pytorch>=0.25.6

base_vit = ViT(
    image_size = 256,
    patch_size = 32,
    num_classes = 1000,
    dim = 512,
    depth = 6,
    heads = 16,
    mlp_dim = 2048,
    dropout = 0.1,
    emb_dropout = 0.1
)

vit = Extractor(
    base_vit,
    return_embeddings_only = True
)

clip = CLIP(
    image_encoder = vit,
    dim_image = 512,           # must match the vision transformer dimension
    dim_text = 512,
    dim_latent = 512,
    num_text_tokens = 10000,
    text_enc_depth = 6,
    text_seq_len = 256,
    text_heads = 8
)

text = torch.randint(0, 10000, (4, 256))
images = torch.randn(4, 3, 256, 256)

loss = clip(text, images, return_loss = True)
loss.backward()