使用Pytorch从零开始实现CLIP

原创

已于 2023-12-05 15:55:06 修改 · 5.2k 阅读

于 2023-11-30 10:23:46 首次发布

本文介绍在PyTorch中从零实现CLIP模型。CLIP可学习句子与图像关系，能检索相关图像，还可作分类器。文中使用Flickr 8k数据集，分别构建图像、文本编码器和Projection Head，实现CLIP模型并训练，最后进行推理，展示模型检索图像的效果。

生成式建模知识回顾:
[1] 生成式建模概述
[2] Transformer I，Transformer II
[3] 变分自编码器
[4] 生成对抗网络，高级生成对抗网络 I，高级生成对抗网络 II
[5] 自回归模型
[6] 归一化流模型
[7] 基于能量的模型
[8] 扩散模型 I, 扩散模型 II

引言

2021 年 1 月，OpenAI 宣布了两种新模型：DALL-E 和 CLIP，这两种模型都是以某种方式连接文本和图像的多模态模型。在本文中，我们将在PyTorch中从零开始实现 CLIP 模型。OpenAI 开源了一些与 CLIP 模型相关的代码，但我发现它令人生畏，而且并不简洁。

CLIP 有什么作用？为什么有趣？

在《Learning Transferable Visual Models From Natural Language Supervision》论文中，OpenAI 介绍了他们的新模型，称为CLIP，用于Contrastive Language-Image Pre-training。简而言之，该模型学习整个句子与其描述的图像之间的关系；从某种意义上说，当训练模型时，给定一个输入句子，它将能够检索与该句子相对应的最相关的图像。这里重要的是，它是在完整的句子上进行训练，而不是像car、dog等单一类别一样。直觉上，当在整个句子上进行训练时，模型可以学习更多的东西，并找到图像和文本之间的一些模式。

他们还表明，当该模型在巨大的图像数据集及其相应文本上进行训练时，它也可以充当分类器。我鼓励你研究论文原文，以更多地了解这个令人兴奋的模型及其在基准数据集上的惊人结果。仅举一例，使用此策略训练的 CLIP 模型对 ImageNet 的分类效果比在 ImageNet 本身上训练的 SOTA 模型更好，该 SOTA 模型专门针对单一分类任务进行了优化！

先跳过过程，让我们看看我们将在本文中从头开始构建的最终模型能够实现什么功能：给出诸如“一个男孩用滑板跳跃”或“一个女孩从秋千上跳跃”这样的查询，模型将检索最相关的图像：
在这里插入图片描述

开始

让我们直接看它的 PyTorch 实现。首先，我们需要一个包含图像和一些描述它们的文本的数据集。坦率地说，网上有很多可用的。我们将使用Flickr 8k 数据集（您可以使用更大的 30k 版本，最终模型的性能会更好），该数据集主要用于图像字幕任务。但是, 我们也可以用它来训练 CLIP 模型。

以下代码将下载 8k（如果取消注释最后几行，则下载 30k）并解压缩它们。Kaggle数据集之下载可参考前文。

!pip install kaggle --upgrade
import os
os.environ['KAGGLE_USERNAME'] = "XXXXXX"
os.environ['KAGGLE_KEY'] = "XXXXXXXXXXXXXXXXXXXXXX" # Enter your Kaggle key here

# For Flickr 8k
!kaggle datasets download -d adityajn105/flickr8k
!unzip flickr8k.zip
dataset = "8k"


# For Flickr 30k
# !kaggle datasets download -d hsankesara/flickr-image-dataset
# !unzip flickr-image-dataset.zip
# dataset = "30k"

关于此数据集需要注意的一件事是: 每张图像都有 5 个标题。后面写损失函数的时候再讲这个！

数据集

正如您在本文的标题图片中看到的，我们需要对图像及其描述文本进行编码。因此，数据集需要返回图像和文本。当然，我们不会将原始文本提供给我们的文本编码器！我们将使用HuggingFace库中的DistilBERT模型（它比 BERT 小，但性能几乎与 BERT 一样）作为我们的文本编码器；因此，我们需要使用 DistilBERT 分词器对句子（标题）进行分词，然后将分词 id (input_ids) 和注意力掩码提供给 DistilBERT。因此，数据集也需要处理标记化。您可以在下面看到数据集的代码。下面我将解释代码中发生的最重要的事情。

关于配置和CFG的说明：我用 python 脚本编写了代码，然后将其转换为 Jupyter Notebook。因此，对于 python 脚本，config 是一个普通的 python 文件，我在其中放置所有超参数，对于 Jupyter Notebook，它是在笔记本开头定义的一个类，用于保留所有超参数。查看GitHub 存储库或笔记本以查看所有超参数。

import os
import cv2
import torch
import albumentations as A

import config as CFG

class CLIPDataset(torch.utils.data.Dataset):
    def __init__(self, image_filenames, captions, tokenizer, transforms):
        """
        image_filenames and cpations must have the same length; so, if there are
        multiple captions for each image, the image_filenames must have repetitive
        file names 
        """

        self.image_filenames = image_filenames
        self.captions = list(captions)
        self.encoded_captions = tokenizer(
            list(captions), padding=True, truncation=True, max_length=CFG.max_length
        )
        self.transforms = transforms

    def __getitem__(self, idx):
        item = {
   
   
            key: torch.tensor(values[idx])
            for key, values in self.encoded_captions.items()
        }

        image = cv2.imread(f"{
     
     CFG.image_path}/{
     
     self.image_filenames[idx]}")
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image = self.transforms(image=image)['image']
        item['image'] = torch.tensor(image).permute(2, 0, 1).float()
        item['caption'] = self.captions[idx]

        return item


    def __len__(self):
        return len(self.captions)



def get_transforms(mode=