contrastive-image-text

VisionTextDualEncoder and CLIP-like model training examples

This folder contains two examples:

1. The first one showcases how to train a CLIP-like vision-text dual encoder model using pre-trained vision and text encoders. The model is inspired by CLIP, introduced by Alec Radford et al. The idea is to train a vision encoder and a text encoder jointly to project the representation of images and their captions into the same embedding space, such that the caption embeddings are located near the embeddings of the images they describe.
2. The second one showcases how to train a BridgeTower model. This model contains bridges between the text and vision encoders that are linked to a cross-modal encoder. This enables effective bottom-up cross-modal alignment between visual and textual representations at different semantic levels in the cross-modal encoder.

Such models can be used for natural language image search and potentially zero-shot image classification.

Requirements

First, you should install the requirements:

pip install -r requirements.txt

Dataset

Recommended (datasets>=4.0.0): use the COCO captions dataset hosted on the Hub. It provides image–caption pairs and does not require trust_remote_code:

import datasets
ds = datasets.load_dataset("regisss/coco_2017", split="train")

This dataset exposes at least the columns image (PIL image) and caption (string). If you prefer local files, you can also use the built-in Datasets imagefolder builder (not a placeholder) to load images/captions from a directory (it typically expects a small CSV/JSON with columns such as image_path and caption).

CLIP-like models

Here is how to run the run_clip.py script for training CLIP-like models.

Create a model from a vision encoder model and a text encoder model

Next, we create a VisionTextDualEncoderModel. The VisionTextDualEncoderModel class lets you load any vision and text encoder model to create a dual encoder. Here is an example of how to load the model using pre-trained vision and text models.

from transformers import (
    VisionTextDualEncoderModel,
    VisionTextDualEncoderProcessor,
    AutoTokenizer,
    AutoImageProcessor
)

model = VisionTextDualEncoderModel.from_vision_text_pretrained(
    "openai/clip-vit-base-patch32", "roberta-base"
)

tokenizer = AutoTokenizer.from_pretrained("roberta-base")
image_processor = AutoImageProcessor.from_pretrained("openai/clip-vit-base-patch32")
processor = VisionTextDualEncoderProcessor(image_processor, tokenizer)

# save the model and processor
model.save_pretrained("clip-roberta")
processor.save_pretrained("clip-roberta")

This loads both the text and vision encoders using pre-trained weights, the projection layers are randomly initialized except for CLIP's vision model. If you use CLIP to initialize the vision model then the vision projection weights are also loaded using the pre-trained weights.

Single-HPU training

Finally, we can run the example script to train the model. Run the following command for single-device training:

python run_clip.py \
    --output_dir ./clip-roberta-finetuned \
    --model_name_or_path ./clip-roberta \
    --dataset_name regisss/coco_2017 \
    --image_column image \
    --caption_column caption \
    --remove_unused_columns=False \
    --do_train  --do_eval \
    --per_device_train_batch_size="512" \
    --per_device_eval_batch_size="64" \
    --learning_rate="5e-5" --warmup_steps="0" --weight_decay 0.1 \
    --overwrite_output_dir \
    --save_strategy epoch \
    --use_habana \
    --gaudi_config_name Habana/clip \
    --throughput_warmup_steps 3 \
    --dataloader_num_workers 16 \
    --sdp_on_bf16 \
    --bf16 \
    --torch_compile_backend=hpu_backend \
    --torch_compile

Multi-HPU training

Run the following command for distributed training:

PT_ENABLE_INT64_SUPPORT=1 \
python3 ../gaudi_spawn.py --world_size 8 --use_mpi run_clip.py \
    --output_dir=/tmp/clip_roberta \
    --model_name_or_path=./clip-roberta \
    --dataset_name regisss/coco_2017 \ 
    --image_column image \
    --caption_column caption \
    --remove_unused_columns=False \
    --do_train --do_eval \
    --mediapipe_dataloader \
    --per_device_train_batch_size="64" \
    --per_device_eval_batch_size="64" \
    --learning_rate="5e-5" --warmup_steps="0" --weight_decay 0.1 \
    --overwrite_output_dir \
    --use_habana \
    --use_lazy_mode=False \
    --gaudi_config_name="Habana/clip" \
    --throughput_warmup_steps=3 \
    --save_strategy="no" \
    --dataloader_num_workers=2 \
    --use_hpu_graphs \
    --max_steps=100 \
    --torch_compile_backend=hpu_backend \
    --torch_compile \
    --logging_nan_inf_filter \

Note

Please note --mediapipe_dataloader is not supported on Gaudi1.

DeepSpeed

You can check the DeepSpeed section in Optimum Habana examples for how to run DeepSpeed. You can also look at the documentation for more information about how to use DeepSpeed in Optimum Habana.

BridgeTower

For training BridgeTower, you need to run the run_bridgetower.py script. For instance, to reproduce the results presented in this blog post, you should run:

PT_HPU_LAZY_MODE=1 python ../gaudi_spawn.py --use_mpi --world_size 8 run_bridgetower.py \
  --output_dir /tmp/bridgetower-test \
  --model_name_or_path BridgeTower/bridgetower-large-itm-mlm-itc \
  --dataset_name jmhessel/newyorker_caption_contest --dataset_config_name matching \
  --image_column image --caption_column image_description \
  --remove_unused_columns=False \
  --do_train --do_eval --do_predict \
  --per_device_train_batch_size="40" --per_device_eval_batch_size="16" \
  --num_train_epochs 5 \
  --learning_rate="1e-5" \
  --overwrite_output_dir \
  --save_strategy no \
  --use_habana --use_lazy_mode --use_hpu_graphs_for_inference --gaudi_config_name Habana/clip \
  --throughput_warmup_steps 3 \
  --logging_steps 10 \
  --dataloader_num_workers 1 \
  --sdp_on_bf16

Note

Please note --mediapipe_dataloader is not supported on Gaudi1.

Inference

To run only inference, you can start from the commands above and you just have to remove the training-only arguments such as --do_train, --per_device_train_batch_size, --num_train_epochs, etc...

For instance, you can run inference with CLIP on COCO on 1 Gaudi card with the following command:

PT_HPU_LAZY_MODE=1 python run_clip.py \
    --output_dir ./clip-roberta-finetuned \
    --model_name_or_path ./clip-roberta \
    --dataset_name regisss/coco_2017 \
    --image_column image \
    --caption_column caption \
    --remove_unused_columns=False \
    --do_eval \
    --per_device_eval_batch_size="64" \
    --overwrite_output_dir \
    --use_habana \
    --use_lazy_mode \
    --use_hpu_graphs_for_inference \
    --gaudi_config_name Habana/clip \
    --bf16 \
    --sdp_on_bf16 \
    --mediapipe_dataloader \

Note

Please note --mediapipe_dataloader is not supported on Gaudi1.