Align before Fuse: Vision and Language Representation Learning with Momentum Distillation (Salesforce Research)
This is the official PyTorch implementation of the ALBEF paper [Blog]. This repository supports pre-training on custom datasets, as well as finetuning on VQA, SNLI-VE, NLVR2, Image-Text Retrieval on MSCOCO and Flickr30k, and visual grounding on RefCOCO+. Pre-trained and finetuned checkpoints are released.
Requirements:
- pytorch 1.8.0
- transformers 4.8.1
- timm 0.4.9
Download:
- Pre-trained checkpoint
- Dataset json files
- Finetuned checkpoint for retrieval on MSCOCO
- Finetuned checkpoint for VQA
- Finetuned checkpoint for visual grounding on RefCOCO+
Visualization:
We provide code in visualize.ipynb to visualize the important areas in an image for each word in a text. Here is an example visualization using the visual grounding checkpoint.
Pre-training on custom datasets:
- Prepare training json files where each json file contains a list. Each item in the list is a dictonary with two key-value pairs: {'image': path_of_image, 'caption': text_of_image}.
- In configs/Pretrain.yaml, set the paths for the json files.
- Pre-train the model using 8 A100 GPUs:
python -m torch.distributed.launch --nproc_per_node=8 --use_env Pretrain.py --config ./configs/Pretrain.yaml --output_dir output/Pretrain
Image-Text Retrieval:
- Download MSCOCO or Flickr30k datasets from the original websites.
- Download and extract the provided dataset json files.
- In configs/Retrieval_coco.yaml or configs/Retrieval_flickr.yaml, set the paths for the json files and the image path.
- Finetune the pre-trained checkpoint using 8 A100 GPUs:
python -m torch.distributed.launch --nproc_per_node=8 --use_env Retrieval.py \ --config ./configs/Retrieval_flickr.yaml \ --output_dir output/Retrieval_flickr \ --checkpoint [Pretrained checkpoint]
VQA:
- Download VQA v2 dataset and Visual Genome dataset from the original websites.
- Download and extract the provided dataset json files.
- In configs/VQA.yaml, set the paths for the json files and the image paths.
- Finetune the pre-trained checkpoint using 8 A100 GPUs:
python -m torch.distributed.launch --nproc_per_node=8 --use_env VQA.py \ --config ./configs/VQA.yaml \ --output_dir output/vqa \ --checkpoint [Pretrained checkpoint]
- Evaluate the result using the official evaluation server.
Visual Entailment:
- Download SNLI-VE dataset from the original website.
- Download and extract the provided dataset json files.
- In configs/VE.yaml, set the paths for the json files and the image path.
- Finetune the pre-trained checkpoint using 8 A100 GPUs:
python -m torch.distributed.launch --nproc_per_node=8 --use_env VE.py \ --config ./configs/VE.yaml \ --output_dir output/VE \ --checkpoint [Pretrained checkpoint]
Visual Grounding on RefCOCO+:
- Download MSCOCO dataset from the original website.
- Download and extract the provided dataset json files.
- In configs/Grounding.yaml, set the paths for the json files and the image path.
- Finetune the pre-trained checkpoint using 8 A100 GPUs:
python -m torch.distributed.launch --nproc_per_node=8 --use_env Grounding.py \ --config ./configs/Grounding.yaml \ --output_dir output/RefCOCO \ --gradcam_mode itm \ --block_num 8 \ --checkpoint [Pretrained checkpoint]
NLVR2:
NLVR2 requires an additional pre-training step with text-assignment (TA) to adapt the model for image-pair inputs. In order to perform TA, first set the paths for the json training files in configs/NLVR_pretrain.yaml, then run:
python -m torch.distributed.launch --nproc_per_node=8 --use_env Pretrain_nlvr.py \ --config ./configs/NLVR_pretrain.yaml \ --output_dir output/NLVR_pretrain \ --checkpoint [Pretrained checkpoint]
We provide the checkpoint after TA pre-training, which can be fine-tuned with the following steps.
- Download NLVR2 dataset from the original website.
- Download and extract the provided dataset json files.
- In configs/NLVR.yaml, set the paths for the json files and the image path.
- Finetune the pre-trained checkpoint using 8 A100 GPUs:
python -m torch.distributed.launch --nproc_per_node=8 --use_env NLVR.py \ --config ./configs/NLVR.yaml \ --output_dir output/NLVR \ --checkpoint [TA pretrained checkpoint]
Citation
If you find this code to be useful for your research, please consider citing.
@article{ALBEF,
title={Align before Fuse: Vision and Language Representation Learning with Momentum Distillation},
author={Junnan Li and Ramprasaath R. Selvaraju and Akhilesh Deepak Gotmare and Shafiq Joty and Caiming Xiong and Steven Hoi},
year={2021},
journal={arXiv preprint arXiv:2107.07651},
}
67 Dec 20, 2022
47 Jan 9, 2023
18 Jan 1, 2023
845 Jan 3, 2023
2 Jan 27, 2022
2 Oct 6, 2021
121 Sep 26, 2022
288 Nov 19, 2022
10 Sep 26, 2022
196 Dec 13, 2022
1 Nov 12, 2021
3 Oct 31, 2022
6 Sep 21, 2022
45 Dec 12, 2022
79 Dec 23, 2022
62 Dec 24, 2022
15 Dec 22, 2022
1.3k Dec 31, 2022
41 Jan 6, 2023