FlagAI进入多模态领域，现已支持一键训练推理CLIP

简介

上一篇我们说道，FlagAI的模型库又有两位成员加入，分别为SwinTransformerV1、和SwinTransformerV2(见FlagAI支持Swin Transformer ），目前在自然语言处理、计算机视觉方向FlagAI已经支持了不少模型，那怎么能缺少多模态方向呢？最近FlagAI也是将Clip一起支持了一下，并开放了训练、推理的例子，欢迎大家尝试！

FlagAI目前陆续支持不同的计算机视觉、自然语言处理、多模态模型，可以方便的支持模型下载，参数加载，训练，推理，无缝切换多GPU训练，deepspeed显存优化，训练加速等全流程操作。尤其针对于加速，显存优化，可以方便的通过设置不同的运行参数即可轻松实现！开源地址为：https://github.com/FlagAI-Open/FlagAI

Clip

Clip可以说是OpenAI在多模态预训练模型的开山之作，首创了利用图文对数据以对比学习的方式进行预训练，重点是方法还简单，让人一目了然，这篇文章就来带大家简单了解一下，详细的技术细节还请阅读原论文：https://arxiv.org/pdf/2103.00020.pdf

模型

Clip模型首先需要两个编码器，分别为图像编码器与文本编码器，具体的模型结构没有太大的要求，目前比较常见的就是都使用Transformer结构即可。

训练

为了训练Clip模型，OpenAI采集了超过4亿的图像-文本对。结合流程图，可以看到，文本模型与图像模型分别对相应的数据进行编码，得到特征，那么一个Batch中对应的文本与图像特征（对角线）肯定是属于同一类（正样本），那么对比学习时便要拉近彼此的距离，反之，一个Batch中不对应的文本与图像特征（非对角线）则属于不同类（负样本），因此要拉远彼此的距离。

推理

当训练好了Clip的两个模型以后，推理的时候便简单多了，比如我们想要识别一张猫的图片，那么我们就可以准备多个语句，例如: ["this is a cat", "this is a dog", "this is a tiger", ...]，利用文本编码器分别计算出这几个语句的特征向量，然后计算当前图像的特征向量，计算图像向量与文本向量之间的相似度，看跟哪个文本的相似度最高，就属于哪个类别即可。

优势

其实Clip的思路不难，不过优势却是非常大的，例如训练高效，标签灵活，可以自由设置语句作为不同类别，并且其鲁棒性非常高，可以看到下图，在多个不同的数据集上，Clip都轻松击败了在ImageNet上进行有监督训练的ResNet101。

扩展

Clip模型自从发布后，许多工作都使用其作为预训练模型，例如DALL-E将其作为排序模型，用来选择所生成的图片中哪个与输入文本相似度最高。视频理解、图像编辑、图像生成等多个任务中，Clip都有极大的用武之地，是一个非常不错的工作。

FlagAI一键调用、加速

代码位置

FlagAI目前已经支持Clip模型进行预训练，推理，样例数据为Imagenet，数据与代码位于FlagAI官方仓库中的examples目录下：

📦examples

┗ 📂clip

┃ ┣ 📂imagenet2012

┃ ┣ 📜deepspeed.json

┃ ┣ 📜hostfile

┃ ┣ 📜inference_clip.py

┃ ┣ 📜train_clip.py

┃ ┗ 📜train_clip_deepspeed.py

其中使用train_clip_deepspeed.py 脚本可以支持多卡deepspeed显存优化技术，多节点训练，更多的参数在deepspeed.json配置文件中修改即可。附上deepspeed工具包的详细配置参数列表：https://www.deepspeed.ai/docs/config-json/

hostfile文件里面为节点的name与ip，用于多节点训练。

训练代码

样例数据位置: "./data/"，其中图文匹配关系在pairs.csv文件中，图像文件在img文件夹中，训练自己的数据集只需要根据数据集特点自定义DataSet即可。

import torch
from flagai.data.dataset.mm.clip_dataset import CsvDataset, clip_transform, collate_fn
from flagai.trainer import Trainer
from flagai.auto_model.auto_loader import AutoLoader

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# cd examples/clip
data_path = "./data/pairs.csv"
img_dir = "./data/img"

trainer = Trainer(env_type="pytorch",
                  epochs=5,
                  pytorch_device=device,
                  batch_size=64,
                  lr=1e-4,
                  log_interval=10,
                  )

loader = AutoLoader(task_name="txt_img_matching",#contrastive learning
                    model_name="clip-base-p32-224",
                    )
model = loader.get_model()
tokenizer = loader.get_tokenizer()

transform = clip_transform(img_size=model.image_size)
train_dataset = CsvDataset(data_path,
                            img_dir,
                            transform,
                            tokenizer)

optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)
trainer.train(model,
              optimizer=optimizer,
              train_dataset=train_dataset,
              collate_fn=collate_fn)

推理代码

以下代码为使用预训练好的模型直接在推理测试图像，预测图像与文本匹配关系。

import torch
from PIL import Image
from flagai.auto_model.auto_loader import AutoLoader
from flagai.data.dataset.mm.clip_dataset import clip_transform

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

loader = AutoLoader(task_name="txt_img_matching", #contrastive learning
                    model_name="clip-base-p32-224")

model = loader.get_model()
model.eval()
model.to(device)
tokenizer = loader.get_tokenizer()
transform = clip_transform(img_size=model.image_size)

def inference():
    image = Image.open("./CLIP.png")
    image = transform(image).unsqueeze(0).to(device)
    text = tokenizer.tokenize_as_tensor(["a diagram", "a dog", "a cat"]).to(device)

    with torch.no_grad():
        image_features = model.encode_image(image)
        text_features = model.encode_text(text)
        text_probs = (image_features @ text_features.T).softmax(dim=-1)

    print(text_probs.cpu().numpy()[0].tolist())

if __name__=="__main__":
    inference()

小结

不难发现，使用FlagAI框架，可以非常轻松的构建好并下载对应模型，切换不同的训练模式进行多GPU训练，加速，并且开始训练。如果想切换为自己的数据集，只需要修改dataset即可。

后续FlagAI会继续支持多模态模型例如diffusion模型，并集成新的任务，敬请期待。

FlagAI特点

FlagAI飞智是一个快速、易于使用和可扩展的AI基础模型工具包。支持一键调用多种主流基础模型，同时适配了中英文多种下游任务。

• FlagAI支持最高百亿参数的悟道GLM(详见GLM介绍)，同时也支持BERT、RoBERTa、GPT2、T5 模型、Meta OPT模型、ViT、SwinTransformer模型和 Huggingface Transformers 的模型。

• FlagAI提供 API 以快速下载并在给定（中/英文）文本上使用这些预训练模型，你可以在自己的数据集上对其进行微调(fine-tuning)或者应用提示学习(prompt-tuning)。

• FlagAI提供丰富的基础模型下游任务支持，例如文本分类、信息提取、问答、摘要、文本生成、图像分类、图文匹配等，对中英文都有很好的支持。

• FlagAI由三个最流行的数据/模型并行库（PyTorch/Deepspeed/Megatron-LM）提供支持，它们之间实现了无缝集成。在FlagAI上，你可以用不到十行代码来并行你的训练、测试过程，也可以方便的使用各种模型提速技巧。

开源项目地址：https://github.com/FlagAI-Open/FlagAI

内容中包含的图片若涉及版权问题，请及时与我们联系删除

FlagAI进入多模态领域，现已支持一键训练推理CLIP

评论