CoCa-PyTorch

repository·main·Indexed 22 days ago

https://github.com/lucidrains/coca-pytorch

An implementation of the Contrastive Captioners (CoCa) model in PyTorch. It combines contrastive learning with an image-to-text transformer decoder to function as an image-text foundation model for captioning and contrastive retrieval. Requires vit-pytorch version 0.40.2 or higher for the image encoder.

Tokens: 1K
Snippets: 5
Records: 5
Agent score: 29%

What's inside coca-pytorch

Install CoCa-PyTorch and dependencies
main
To use this project, you need to install coca-pytorch and a compatible version of vit-pytorch (version 0.40.2 or higher) to provide the image encoder.
```
$ pip install coca-pytorch
$ pip install vit-pytorch>=0.40.2
```
https://github.com/lucidrains/coca-pytorch
Train CoCa with caption and contrastive loss
main
To train the model, pass text and images tensors to the coca instance and set return_loss = True. This returns the combined loss from both the autoregressive captioning task and the contrastive image-text task.
```
# text shape: (batch, seq_len)
# images shape: (batch, channels, height, width)

loss = coca(
    text = text,
    images = images,
    return_loss = True
)

loss.backward()
```
https://github.com/lucidrains/coca-pytorch
Get CLIP-like text and image embeddings
main
To obtain the contrastive embeddings (similar to CLIP) for both text and images, call the model with return_embeddings = True.
```
# Returns text_embeds and image_embeds, both of shape (batch, dim)
text_embeds, image_embeds = coca(
    text = text,
    images = images,
    return_embeddings = True
)
```
https://github.com/lucidrains/coca-pytorch
Get caption logits from CoCa
main
To perform captioning, call the model with text and images. This returns the logits for the text tokens.
```
# Returns logits of shape (batch, seq_len, num_tokens)
logits = coca(
    text = text,
    images = images
)
```
https://github.com/lucidrains/coca-pytorch
Instantiate the CoCa model
main
The CoCa class implements the Contrastive Captioners architecture. It requires a vision transformer (ViT) as the img_encoder.
Important: The vision transformer must be wrapped in an Extractor with return_embeddings_only = True so that it returns embeddings of shape (batch, seq, dim) instead of class logits.
Key parameters for CoCa:
- dim: Model dimension.
- img_encoder: The vision transformer (wrapped in Extractor).
- image_dim: Dimension of the image embeddings (if different from dim).
- num_tokens: Vocabulary size for text tokens.
- unimodal_depth: Depth of the unimodal transformer.
- multimodal_depth: Depth of the multimodal transformer.
- dim_head: Dimension per attention head.
- heads: Number of attention heads.
- caption_loss_weight: Weight for the autoregressive caption loss.
- contrastive_loss_weight: Weight for the contrastive loss between image and text CLS embeddings.
```
from vit_pytorch.simple_vit_with_patch_dropout import SimpleViT
from vit_pytorch.extractor import Extractor
from coca_pytorch.coca_pytorch import CoCa
import torch

# 1. Setup the vision transformer with Extractor
vit = SimpleViT(
    image_size = 256,
    patch_size = 32,
    num_classes = 1000,
    dim = 1024,
    depth = 6,
    heads = 16,
    mlp_dim = 2048,
    patch_dropout = 0.5
)
vit = Extractor(vit, return_embeddings_only = True, detach = False)

# 2. Instantiate CoCa
coca = CoCa(
    dim = 512,
    img_encoder = vit,
    image_dim = 1024,
    num_tokens = 20000,
    unimodal_depth = 6,
    multimodal_depth = 6,
    dim_head = 64,
    heads = 8,
    caption_loss_weight = 1.,
    contrastive_loss_weight = 1.
).cuda()
```
https://github.com/lucidrains/coca-pytorch