史上最大CLIP模型ViT-G/14開源：ImageNet可達(dá)80.1

LAION這個(gè)開源機(jī)構(gòu)相信大家都不陌生，這個(gè)非盈利開源機(jī)構(gòu)先后開源了LAION-400M和LAION-5B等大規(guī)模圖文對(duì)數(shù)據(jù)集，而且也發(fā)起了OpenCLIP項(xiàng)目，在22年12月份LAION也發(fā)布了基于OpenCLIP的scaling laws報(bào)告，其中最大模型ViT-H/14在ImageNet上的zero-shot準(zhǔn)確度達(dá)到了78.0%。在今年1月底，LAION繼續(xù)擴(kuò)增模型，訓(xùn)練了參數(shù)量約2.5B的ViT-G/14模型，其在ImageNet上的zero-shot準(zhǔn)確度達(dá)到了80.1%，這也是目前開源的最大的CLIP模型。相比之前的OpenCLIP H/14模型，新開源的OpenCLIP G/14模型的text encoder參數(shù)從原來的354.0M增加到694.7M，而image encoder參數(shù)從632.08M增加至1844.9M，增加接近3倍。OpenCLIP G/14模型的加入也可以進(jìn)一步驗(yàn)證scaling law：OpenCLIP G/14模型采用了meta所提出的FLIP中的策略來進(jìn)行訓(xùn)練，F(xiàn)LIP的主要思路是mask一部分patch來訓(xùn)練image encoder，這帶來的好處是減少顯存從而增大batch size。基于FLIP中的patch dropout/mask策略，OpenCLIP G/14模型的batch size達(dá)到了160K（512～760 A100上），batch size對(duì)CLIP的訓(xùn)練效果是至關(guān)重要的。具體來說，OpenCLIP G/14模型的訓(xùn)練共分為兩個(gè)階段（遵循FLIP）：

第一階段：Patch dropout

采用50%的patch dropout在LAION-2B數(shù)據(jù)集上訓(xùn)練，共采樣32B的樣本量（訓(xùn)練所過的全部樣本數(shù)量）。這里采用的batch size是160K，優(yōu)化器為AdamW，學(xué)習(xí)速率為2e-3，采用cosine decay schedule。這個(gè)階段后，模型可以在ImageNet上達(dá)到79.07。

第二階段：Unmasked tuning + Model soups

patch mask或dropout策略會(huì)導(dǎo)致image encoder在訓(xùn)練和測(cè)試時(shí)的不一致，所以需要少量的unmasked tuning，此時(shí)就和常規(guī)的CLIP一樣不進(jìn)行mask，注意這個(gè)階段需要采用gradient checkpointing來保證160K的batch size。這個(gè)階段共采用3種的不同的設(shè)置（學(xué)習(xí)速率，訓(xùn)練樣本等設(shè)置不同）來訓(xùn)練3個(gè)不同的模型，其分別可以在ImageNet上達(dá)到79.43、79.45和79.2。最后再基于谷歌論文Model soups中提取的權(quán)重平均來得到最后的模型，其在ImageNet上達(dá)到80.1。

下圖為OpenCLIP G/14模型和OpenCLIP H/14模型在各個(gè)圖像分類數(shù)據(jù)集上的zero-shot性能對(duì)比，可以看到OpenCLIP G/14模型在絕大多數(shù)數(shù)據(jù)集上均明顯優(yōu)于OpenCLIP H/14模型。

CLIP模型的應(yīng)用甚廣，除了常規(guī)圖像和文本多模態(tài)檢索任務(wù)，CLIP還可以應(yīng)用在文生圖大模型上，比如DALLE2和Stable Diffusion均采用CLIP模型所提取的特征作為擴(kuò)散模型的輸入條件。近期所發(fā)布的Stable Diffusion 2.0相比1.5版本，一個(gè)最重要的改進(jìn)就是將text encoder從原來的ViT-L/14換成了更大的ViT-H/14（參數(shù)量增大3倍），SD 2.0相比1.5版本在FID和CLIP score上均有一定提升。毫無疑問，更大的模型ViT-G/14會(huì)帶來進(jìn)一步的提升。

目前OpenCLIP G/14模型已經(jīng)在hugging face上開源：https://huggingface.co/laion/CLIP-ViT-bigG-14-laion2B-39B-b160k。你可以直接基于openclip庫來使用它：

import open_clip
import torch
from PIL import Image

model, preprocess_train, preprocess_val = open_clip.create_model_and_transforms('hf-hub:laion/CLIP-ViT-bigG-14-laion2B-39B-b160k')
tokenizer = open_clip.get_tokenizer('hf-hub:laion/CLIP-ViT-bigG-14-laion2B-39B-b160k')

image = preprocess_val(Image.open("CLIP.png")).unsqueeze(0)
text = tokenizer(["a diagram", "a dog", "a cat"])

with torch.no_grad(), torch.cuda.amp.autocast():
    image_features = model.encode_image(image)
    text_features = model.encode_text(text)
    image_features /= image_features.norm(dim=-1, keepdim=True)
    text_features /= text_features.norm(dim=-1, keepdim=True)

    text_probs = (100.0 * image_features @ text_features.T).softmax(dim=-1)

print("Label probs:", text_probs)  # prints: [[1., 0., 0.]]

參考

• Reproducible scaling laws for contrastive language-image learning
• https://github.com/mlfoundations/open_clip