Multilingual CLIP is a pre-trained model which can be used for multilingual semantic search and zero-shot image classification in 100 languages.
Model Architecture
Multilingual CLIP was built using OpenAI CLIP model. I have used the same Vision encoder (ResNet 50x4), but instead I replaced their text encoder (Transformer) with a Mulilingual Text Encoder (XLM-Roberta) and a configurable number of projection heads, as seen below:
Model Architecture
The model was trained in a distributed fashion on 16 Habana Gaudi Accelerators and with mixed Precision in two phases (using COCO Dataset for phase 1 and Google Conceptual Captions for phase 2). The training pipeline was built using PyTorch, PyTorch Lightning, and Distributed Data Parallel.
Datasets
Three datasets have been used for building the model. COCO captions was used for training phase 1 and Google Conceptual Captions was used for training phase 2. Unsplash dataset was used for testing and inference.
COCO Captions
COCO (Common Objects in Context) is a large-scale object detection, segmentation, and captioning dataset. The COCO captions dataset has around ~85000 images and captions pairs.
Run the following to download the dataset:
./download_coco.sh
This dataset was used for the first pre-training phase.
Google Conceptual Captions
Conceptual Captions is a dataset consisting of ~3.3 million images annotated with captions. In contrast with the curated style of other image caption annotations, Conceptual Caption images and their raw descriptions are harvested from the web, and therefore represent a wider variety of styles.
Download the datasets urls/captions from here as save it to datasets/googlecc/googlecc.tsv. The full dataset has over 3 million images, but you can select a subset by loading the googlecc.tsv file and saving only the number of rows you want (I have used 1 million images for training).
Then run the following commands to download each image on the googlecc.tsv file:
bash
1npminstall2node download_build_googlecc.js
This dataset was used for the second pre-training phase.
Unplash
This dataset was used as the test set during inference.
Run python3.8 download_unsplash.py to download the dataset.
Copy id_rsa.pub contents from every node's docker to every other node's docker's ~/.ssh/authorized_keys (all public keys need to be in all hosts' authorized_keys):
--dataset-num-workers Number of workers (default: 8)
--dataset-type Dataset type (coco or googlecc) (default: coco)
--dataset-dir Dataset dir (default: ./datasets/coco/)
--dataset-subset-size Load only a subset of the dataset (useful for debugging)
--dataset-train-split Dataset train split (default: 0.8)
--train-device Type of device to use (default: hpu)
--distributed-num-nodes Number of nodes (machines) (default: 2)
--distributed-parallel-devices Number of parallel devices per node (default: 8)
--distributed-master-address Master node IP address
--distributed-master-port Master node port (default: 12345)
--distributed-bucket-cap-mb DDP bucket cap MB (default: 200)
--checkpoint-dir Model checkpoint dir (default: ./models)
--checkpoint-save-every-n Save every n epochs (default: 1)
--checkpoint-load-vision-path Load vision encoder checkpoint
--checkpoint-load-text-path Load text encoder checkpoint
--model-visual-name Which visual model to use (default: RN50x4)
--model-textual-name Which textual model to use (default: xlm-roberta-base)
--hyperparam-num-layers Number of layers (default: 3)
--hyperparam-lr Model learning rate (default: 0.001)
--hyperparam-epochs Max epochs (default: 100)
--hyperparam-precision Precision (default: 16)
--hyperparam-batch-size Batch size (default: 64)
--wandb-project W&B project name (default: clip)
--wandb-enabled W&B is enabled? (default: True)
1from models import MultiLingualCLIP, load_model
23text_checkpoint_path ='/path/to/text model checkpoint'4vision_checkpoint_path ='/path/to/vision model checkpoint'56model = MultiLingualCLIP(num_layers=3)7load_model(model, vision_checkpoint_path, text_checkpoint_path)
Generate embeddings
Run the following (after downloading Unplash dataset):
python3.8 ./generate_embeddings.py
Searching images
python
1import numpy as np
2from search import MultiLingualSearch
34images_embeddings = np.load('/path/to/images_embeddings')5images_data =[...]# List of image info for each row of the embeddings. For instance, it could be a list of urls, filepaths, ids. They will be returned when calling the search function6semantic_search = MultiLingualSearch(model, images_embeddings, images_data)78results = semantic_search.search('विद्यालय में')# Means at school9print(results)
