Package: clip-text-decoder
clip-text-decoderTrain an image captioner with 0.323 BLEU on COCO Captions in under one hour! (0.352 BLEU with beam search🙂)Generates text captions for images from their embeddings.Now includesBLIPas an available vision backbone!Example PredictionsComputed using the pretrained model mentioned below."A man riding a wave on top of a surfboard.""A baseball player is swinging a bat at a ball.""A dog jumping in the air to catch a frisbee."InstallationUsingpip:pipinstall"clip @ git+https://github.com/openai/CLIP.git"pipinstall"lavis @ git+https://github.com/salesforce/LAVIS.git"pipinstallclip-text-decoderFrom source:pipinstall"clip @ git+https://github.com/openai/CLIP.git"pipinstall"lavis @ git+https://github.com/salesforce/LAVIS.git"gitclonehttps://github.com/fkodom/clip-text-decoder.gitcdclip-text-decoder
pipinstall.InferencePretrained ModelfromPILimportImageimporttorchfromclip_text_decoder.modelimportImageCaptionInferenceModelmodel=ImageCaptionInferenceModel.download_pretrained()device="cuda"iftorch.cuda.is_available()else"cpu"model.to(device)image=Image.open("path/to/image.jpeg")# The beam_size argument is optional. Larger beam_size is slower, but has# slightly higher accuracy. Recommend using beam_size <= 3.caption=model(image,beam_size=1)To cache the pretrained model locally, so that it's not re-downloaded each time:model=ImageCaptionInferenceModel.download_pretrained("path/to/model.pt")Custom Trained ModelTraining produces amodel.ptarchive, containing aTokenizerand model parameters.  To reload the trained inference model:fromclip_text_decoder.modelimportImageCaptionInferenceModelmodel=ImageCaptionInferenceModel.load("path/to/model.pt").to(device)# Load image and get predictions like above...Ablation: Beam SizeMeasuring the BLEU-4 score for differentbeam_sizearguments.  By default, the inference model uses a beam size of 1:fromclip_text_decoder.modelimportImageCaptionInferenceModelmodel=ImageCaptionInferenceModel.load("path/to/model.pt")caption=model(image,beam_size=1)Using largerbeam_sizegives better BLEU score with a trade-off of slower inference speeds. The metrics below were collected from the same model, which uses a BLIP vision backbone and was trained for 10 epochs (roughly 1 hour on a T4 GPU):Beam sizeBLEU-41 (default)0.32320.34330.35040.352TrainingLaunch your own training session usingtrain.py:pythontrain.py--max-epochs10Training CLI arguments, along with their default values:--vision-backboneblip:base# (str)--language-modeldistilgpt2# (str)--max-epochs10# (int)--beam-size1# (int)--batch-size32# (int)--accumulate-grad-batches4# (int)--precision16# (16 or 32)--seed0# (int)One epoch takes about 5-6 minutes using a T4 GPU, which is usually free in Google Colab (depending on availability).  After about 10 training epochs, you'll reach a BLEU-4 score just over 0.30 (without beam search).  So, in under an hour, you can train a pretty good image captioning model. 😎NotesBLEU doesn't increase much beyond 1 hour of training. Training and validation loss will continue to decrease, but the resulting image captions are effectively equivalent.This appears to be a limitation of the image embeddings, rather than a limitation of the language model. Changing the vision backbone gives the biggest improvement in BLEU score. (BLIP gets 5-10% better BLEU than CLIP backbones using the same language model head.) Larger language models (e.g. GPT-2 Large) don't improve the BLEU score by much.TODOPlan to train on Conceptual Captions for more generic image captioning.