学习笔记 – CoCoOp代码逐行解读

阅读CoCoOp之前,对CLIP的源码没有过多熟悉,所以本篇也是边读边学,尽可能的让本篇代码解读层次递进的包含Transformer,CLIP,以及CoCoOp的主要代码实现,说明关联的原理。若有不清晰之处,随时修订。尔后基于这一Dassl训练框架的CLIP相关实现的解读,就不会再次包含CLIP和Transformer了,避免内容过于繁杂。本篇的目的是清晰理解每一行代码,避免后续工作遇到问题。

Config

本次解读使用的trainer是CoCoOp的vit_b16_c4_ep10_batch1_ctxv1

class CoCoOp(TrainerX)

CoCoOp主训练代码部分

代码有点多,切一部分先

@TRAINER_REGISTRY.register()
class CoCoOp(TrainerX):
    def check_cfg(self, cfg):
        assert cfg.TRAINER.COCOOP.PREC in ["fp16", "fp32", "amp"]

这个部分主要是训练器的定义,其中涉及到各种参数定义,以及方法调用

首先是训练器的定义,这个涉及到Dassl.Pytorch这个工具箱,其中@TRAINER_REGISTRY.register()声明注册了新的训练器,而继承的TrainerX代表这是一个域泛化的训练器[1]

check_cfg方法设定了三种精度,全精度fp32,半精度fp16,混合精度AMP,我们的config当中指定的是fp16

接下来是模型的构建部分

    def build_model(self):
        cfg = self.cfg
        classnames = self.dm.dataset.classnames

        print(f"Loading CLIP (backbone: {cfg.MODEL.BACKBONE.NAME})")
        clip_model = load_clip_to_cpu(cfg)

第二第三行分别是配置文件和训练集类名

接下来第六行载入CLIP模型,可以看看load_clip_to_cpu是如何载入的

def load_clip_to_cpu(cfg):
    backbone_name = cfg.MODEL.BACKBONE.NAME
    url = clip._MODELS[backbone_name]
    model_path = clip._download(url)

    try:
        # loading JIT archive
        model = torch.jit.load(model_path, map_location="cpu").eval()
        state_dict = None

    except RuntimeError:
        state_dict = torch.load(model_path, map_location="cpu")

    model = clip.build_model(state_dict or model.state_dict())

    return model

骨干网络名称来自配置文件的定义,我们的配置文件定义的就是“ViT-B/16

模型权重路径是来源于CLIP文件夹下定义的下载路径,感兴趣的可以去看下CLIP源码,分别定义了几种不同的骨干网模型权重的下载地址,这里不赘述。

接下来就是载入权重,然后返回模型。

继续回到构建部分的代码,往下读

        if cfg.TRAINER.COCOOP.PREC == "fp32" or cfg.TRAINER.COCOOP.PREC == "amp":
            # CLIP's default precision is fp16
            clip_model.float()

        print("Building custom CLIP")
        self.model = CustomCLIP(cfg, classnames, clip_model)

因为配置文件已经定义了fp16,所以第一行到第三行可以忽略

接下来进入了自定义模型的构建,其实就是CoCoOp修改CLIP源码的部分了。

我们可以进入到CustomCLIP这个类当中详细看一下

class CustomCLIP(nn.Module)

初始化部分

    def __init__(self, cfg, classnames, clip_model):
        super().__init__()
        self.prompt_learner = PromptLearner(cfg, classnames, clip_model)
        self.tokenized_prompts = self.prompt_learner.tokenized_prompts
        self.image_encoder = clip_model.visual
        self.text_encoder = TextEncoder(clip_model)
        self.logit_scale = clip_model.logit_scale
        self.dtype = clip_model.dtype

首先是self.prompt_learner,将配置,训练集类名,模型输入,构建CoOp的可学习提示词,主要是可以得到一个可学习的参数ctx,也即上下文向量

源论文 Fig. 2示意图

self.prompt_learner当中会对提示词token化,将这些代表了提示词的token提取出来,就是self.tokenized_prompts

所以这两行做的事情其实是一件事情,我们可以进入PromptLearner,看看具体做了什么


class PromptLearner(nn.Module)

CoCoOp的大核心:提示词学习

这个模块的代码蛮长的,我们切成小块一点点看。

class PromptLearner(nn.Module):
    def __init__(self, cfg, classnames, clip_model):
        super().__init__()
        n_cls = len(classnames)
        n_ctx = cfg.TRAINER.COCOOP.N_CTX
        ctx_init = cfg.TRAINER.COCOOP.CTX_INIT
        dtype = clip_model.dtype
        ctx_dim = clip_model.ln_final.weight.shape[0]
        vis_dim = clip_model.visual.output_dim
        clip_imsize = clip_model.visual.input_resolution
        cfg_imsize = cfg.INPUT.SIZE[0]
        assert cfg_imsize == clip_imsize, f"cfg_imsize ({cfg_imsize}) must equal to clip_imsize ({clip_imsize})"

从第四行开始,首先定义n_cls,表示分类数量。值得注意的是,CoCoOp的CLS不是VIT当中的CLS,并不是用于表达内容的,而是用于对齐CLIP的输入的起始符。

第五行是ctx的数量,表达可学习的上下文向量token数量,本处是4,因为我们的训练配置文件是vit_b16_c4_ep10_batch1_ctxv1,其中c4代表ctx的token数为4,可以在配置文件的TRAINER部分看到n_ctx的配置设定,[CTX]的数量决定了动态学习的prompt的长度。

第六行是预定义的ctx,在配置文件中,预定义的ctx是:”a photo of a

第七行为浮点数类型,fp16

第八行定义了ctx的特征维度,为了对齐文本,提示上下文向量的特征维度是由文本特征维度决定的,而文本特征维度是由最后一层归一化层决定的,本处归一化层维度为512。经过归一化后还会经过一层文本映射层,文本维度最终会降维到512。

本处的ctx的特征维度为512,这里不是被文本映射层映射到了512维度,而是文本编码器输出的原始维度就是512,通过检查clip_model.ln_final.weight.shape即可得知其维度为(512,),通过clip_model.text_projection.shape,我们可以得知,文本映射层的结果是(512,512),所以是512 -> 512。CLIP文本编码器的过程和视觉编码器是不同的,因为视觉编码器的特征维度通常是768或者1024,这是由CLIP/ViT的设定决定的。

当然,如果文本特征维度是768/1024,那么ctx就会是768/1024维度,然后会是映射到512层:(768,512),(1024,512)。

第九行为视觉编码器输出的特征维度,本处为标准的ViT-b/16定义的768维,具体细节可以参考第八行。

第十行是输入图像的尺寸,默认为224,所以这里是224。

十一行是我们在配置文件里定义的尺寸,(224,224)取width,所以这里也是224。

第十二行是作一次校验,强制模型输入尺寸和配置相同。224和224尺寸一致,所以通过。

我们继续往下看,接下来是上下文向量的初始化设定,先前提到我们在配置文件中定义了ctx初始化的prompt为:”a photo of a

if ctx_init:
            # use given words to initialize context vectors
            ctx_init = ctx_init.replace("_", " ")
            n_ctx = len(ctx_init.split(" "))
            prompt = clip.tokenize(ctx_init)
            with torch.no_grad():
                embedding = clip_model.token_embedding(prompt).type(dtype)
            ctx_vectors = embedding[0, 1 : 1 + n_ctx, :]
            prompt_prefix = ctx_init
        else:
            # random initialization
            ctx_vectors = torch.empty(n_ctx, ctx_dim, dtype=dtype)
            nn.init.normal_(ctx_vectors, std=0.02)
            prompt_prefix = " ".join(["X"] * n_ctx)

所以在这段代码中,我们会直接执行上面一部分,而下面一部分是随机初始化,意味着其实不指定prompt给CoCoOp的话,模型会直接初始化一个随机的prompt,然后在训练的过程中学习得到一个合理的prompt。值得一提的是,除了这次用于解析的配置文件vit_b16_c4_ep10_batch1_ctxv1是带有预设的prompt之外,其他的都是没有预设的,会进行随机初始化。

考虑到我们的配置文件已经有初始化的ctx了,所以我们主要分析初始化了的部分,第三行为字符串匹配替换,可能是有用“_”作为分割的,但本处初始化的提示词直接就是以空格作为分割的。

第四行为空格作为分割,计算提示词数量,”a photo of a“为4。

第五行是将提示词token化,形状为(1,77),77代表的是CLIP预定义好的上下文长度,所以如果观察上下文向量的话,会发现长度也是77。需要注意的是,这里的prompt是固定的,而不是可学习的上下文向量,不要弄混了,另外假如是随机初始化的话,是没有这一步的。

第六和第七行紧接着将token变换为文本嵌入,嵌入化是由CLIP的词汇表和词向量机制实现的,具体来说,CLIP维护了一个49408长度的词汇表,其中49406代表[CLS]/[SOS],也就是起始符,而49407就是终止符,在49405之前的会作为用于映射的Token ID,组成一组Token,也就是第五行定义的prompt。

而每一个Token都会被词向量机制转换为一组512维度的嵌入向量,512维度用于对齐transformer文本编码器的输入,关于为什么是512维度,后续到文本编码器部分会有说说明。

经过转换后,第七行输出的形状就变为了(1,77,512)

第八行是取出可学习的上下文向量部分,具体来分析下embedding[0, 1 : 1 + n_ctx, :]

根据embedding的形状就可以知道,其只有一个维度,所以第一个为0

进入该维度后,数据为(77,512),其中77表示的是总上下文向量长度,已知第一位为[CLS]/[SOS],中间有四位有效位,而后紧接着是结束符[EOS],所以第一位和有效位后一位是不代表任何意义的。

[1:1+n_ctx],其中1:表示第二位开始,n_ctx的值是4,代表我们有4位有效位,是包含了提示内容的,那么1+n_ctx就是5,1:1+n_ctx就是跳过第零位,从第一位开始,一直读取到第五位(不含第五位)为止,一共四位数,正好是有含义的部分。

最后定义了一个prompt_prefix,是原始未经过处理的提示符字符串。

继续往下看

        self.ctx = nn.Parameter(ctx_vectors)

        self.meta_net = nn.Sequential(OrderedDict([
            ("linear1", nn.Linear(vis_dim, vis_dim // 16)),
            ("relu", nn.ReLU(inplace=True)),
            ("linear2", nn.Linear(vis_dim // 16, ctx_dim))
        ]))
        
        if cfg.TRAINER.COCOOP.PREC == "fp16":
            self.meta_net.half()

        classnames = [name.replace("_", " ") for name in classnames]
        name_lens = [len(_tokenizer.encode(name)) for name in classnames]
        prompts = [prompt_prefix + " " + name + "." for name in classnames]
        tokenized_prompts = torch.cat([clip.tokenize(p) for p in prompts])  # (n_cls, n_tkn)
        with torch.no_grad():
            embedding = clip_model.token_embedding(tokenized_prompts).type(dtype)

第一行让先前初始化好的,仅包含了提示信息且嵌入化后上下文向量变为可学习的参数

接下来第三行定义了一个小型的MLP网络,后续将作用于视觉特征提取,如果配置文件指定半精度则网络适配为半精度。

接下来和先前初始化提示词的过程如出一辙,不过这里的prompts不再是原先固定的”a photo of a“,而是拼接上了类名的prompts,例如:”a photo of a toucan.

第十五行对提示词Token化,形状是(1000,77),是不是很眼熟?1000代表类的数量,77的含义先前说明过了,这里也没有区别。

第十六七行进行了嵌入化,具体原理已赘述过。

        # These token vectors will be saved when in save_model(),
        # but they should be ignored in load_model() as we want to use
        # those computed using the current class names
        self.register_buffer("token_prefix", embedding[:, :1, :])  # SOS
        self.register_buffer("token_suffix", embedding[:, 1 + n_ctx :, :])  # CLS, EOS

        self.n_cls = n_cls
        self.n_ctx = n_ctx
        self.tokenized_prompts = tokenized_prompts  # torch.Tensor
        self.name_lens = name_lens

第四行和第五行是将[SOS]和[EOS]作为参数进行了保存,在保存模型和载入模型时会引入,参数不会受到模型优化的影响。具体可以参考Pytorch的register_buffer的定义[2]

接下来也只是声明一下变量到self当中,方便后续方法读取,初始化函数部分算是结束了。

喘口气,接下来看看前向传播部分:

    def forward(self, im_features):
        prefix = self.token_prefix
        suffix = self.token_suffix
        ctx = self.ctx                     # (n_ctx, ctx_dim)
        bias = self.meta_net(im_features)  # (batch, ctx_dim)
        bias = bias.unsqueeze(1)           # (batch, 1, ctx_dim)
        ctx = ctx.unsqueeze(0)             # (1, n_ctx, ctx_dim)
        ctx_shifted = ctx + bias           # (batch, n_ctx, ctx_dim)
        
        # Use instance-conditioned context tokens for all classes
        prompts = []
        for ctx_shifted_i in ctx_shifted:
            ctx_i = ctx_shifted_i.unsqueeze(0).expand(self.n_cls, -1, -1)
            pts_i = self.construct_prompts(ctx_i, prefix, suffix)  # (n_cls, n_tkn, ctx_dim)
            prompts.append(pts_i)
        prompts = torch.stack(prompts)
        
        return prompts

prefix和suffix分别是[SOS]和[EOS]的Token

ctx是由上下文向量数量和上下文向量维度组成的上下文向量,形状是(4,512),具体是什么先前说过了。

第五行定义的偏移量比较有意思,首先将图像特征进入到之前定义的MLP网络中,这个图像特征的形状是(1,512),是来源于CLIP图像编码器的输出。通过MLP计算出一个偏近于图像的特征值,设定为偏移量bias,最后在第六第七行与上下文向量对齐后通过广播机制加在一起,生成出新的prompt。接下来在训练过程中,反向传播会让meta_net学习到图像特征偏移量与prompt之间的关系,优化得到最佳的prompt,使新的prompt会更贴近输入图像的特征。

关于prompts内容的构成,可以看到第十四行的定义,看向construct_prompts方法:

    def construct_prompts(self, ctx, prefix, suffix, label=None):
        # dim0 is either batch_size (during training) or n_cls (during testing)
        # ctx: context tokens, with shape of (dim0, n_ctx, ctx_dim)
        # prefix: the sos token, with shape of (n_cls, 1, ctx_dim)
        # suffix: remaining tokens, with shape of (n_cls, *, ctx_dim)

        if label is not None:
            prefix = prefix[label]
            suffix = suffix[label]

        prompts = torch.cat(
            [
                prefix,  # (dim0, 1, dim)
                ctx,     # (dim0, n_ctx, dim)
                suffix,  # (dim0, *, dim)
            ],
            dim=1,
        )

        return prompts

可以清晰的看到,prompts就是由prefix,ctx,suffix拼接而成,各维度一致于dim0,prefix就是[SOS],而ctx就是形状为(4,512)的上下文向量,suffix是(72,512),也就是上下文向量之后的[EOS]标签和填充项,一共是1+4+72,刚好77位。


接下来回到自定义CLIP类下,继续看除了可学习的prompt外,CoCoOp还做了什么改动。

self.image_encoder直接调用CLIP的视觉编码器,这个部分是直接调用CLIP模型的编码器,没有什么改动,所以直接粗略看输入输出就好


clip_model.visual

简单说明下CLIP编码器输入输出结构

输入

输入图像为标准的(3,224,224)

CLIP遵从VIT,不同的地方是本处实现用的是卷积层,而不是线性层。

首先将输入图像切成16\times 16大小的Patch,其中会加上RGB通道数,所以每个Patch的形状就是16\times 16\times 3 = 768,再经过映射层。其他大小的ViT基于此,经过线性映射层的时候会将形状变更成1024,2048等。

输入图像形状为224\times 224,按照16\times 16切割成Patch,得到\frac{224\times 224}{16\times 16}=196个Patch。

最后形状为(196,768),首先进入归一化LayerNorm层,最后输入到Transformer编码器中

源论文 Fig. 1 ViT Overview

结合上batch size,那么输入形状就是(1,3,224,224)->(1,196,768)

输出

CLIP输出的是图像特征,输出维度与文本特征(512dim)对齐,所以形状是(1,512)


self.text_encoder = TextEncoder(clip_model) 是文本编码器部分,CLIP利用预训练的Transformer文本编码器提取文本特征。


class TextEncoder(nn.Module)

我们也可以进来看一下此处定义的文本编码器

初始化部分可以一笔带过,主要是继承了CLIP模型当中的各个所需模块

class TextEncoder(nn.Module):
def __init__(self, clip_model):
super().__init__()
self.transformer = clip_model.transformer
self.positional_embedding = clip_model.positional_embedding
self.ln_final = clip_model.ln_final
self.text_projection = clip_model.text_projection
self.dtype = clip_model.dtype

具体看前向传播部分:

    def forward(self, prompts, tokenized_prompts):
        x = prompts + self.positional_embedding.type(self.dtype)
        x = x.permute(1, 0, 2)  # NLD -> LND
        x = self.transformer(x)
        x = x.permute(1, 0, 2)  # LND -> NLD
        x = self.ln_final(x).type(self.dtype)

        # x.shape = [batch_size, n_ctx, transformer.width]
        # take features from the eot embedding (eot_token is the highest number in each sequence)
        x = x[torch.arange(x.shape[0]), tokenized_prompts.argmax(dim=-1)] @ self.text_projection

        return x

其中,第二行是定义输入,prompts是传入的提示词嵌入,也就是之前的tokenized_prompts当中的内容,接下来广播加上可学习的位置编码,实现token顺序化,这个部分是遵循ViT的实现,因为ViT是要求显性的位置编码的。

第三行到第五行,主要是转换格式为Transformer的输入格式,再转换输出回原先输入格式。在输入编码器时,形状是:(batch_size, sequence_length,dim),而Pytorch实现的Transformer所需要的形状是:(sequence_length, batch_size, dim),具体可以看torch.nn.Transformer的实现。[3]

第六行输出进入LayerNorm层进行归一化。

第七行提取嵌入当中最后一个结束符[EOS],该Token包含文本全局的特征信息,是CLIP的实现。

看下来可以得知,CoOp主要引入的就是第一行定义的可学习prompts,其他部分处理与CLIP的文本编码器一致。


class CustomCLIP(nn.Module)

初始化部分

继续回到初始化部分,看向self.logit_scale

翻阅CLIP模型定义可以得知,这是一个CLIP的可训练权重参数,在训练过程中,可以放大图像文本对比学习过程中的相似度,加快收敛速度。[4]对于其实际的实验效果,可以查看引用。

dtype是设定的半精度fp16,不解释。

接下来重新定义了CLIP的前向传播过程:

    def forward(self, image, label=None):
        tokenized_prompts = self.tokenized_prompts
        logit_scale = self.logit_scale.exp()

        image_features = self.image_encoder(image.type(self.dtype))
        image_features = image_features / image_features.norm(dim=-1, keepdim=True)

        prompts = self.prompt_learner(image_features)
        
        logits = []
        for pts_i, imf_i in zip(prompts, image_features):
            text_features = self.text_encoder(pts_i, tokenized_prompts)
            text_features = text_features / text_features.norm(dim=-1, keepdim=True)
            l_i = logit_scale * imf_i @ text_features.t()
            logits.append(l_i)
        logits = torch.stack(logits)
        
        if self.prompt_learner.training:
            return F.cross_entropy(logits, label)
        
        return logits

第二行和第三行分别是token化后的提示和放大因子

接下来图像输入视觉编码器,得到图像特征,再对其作归一化,这里的是CLIP的标准流程。

归一化这里可以拆成两个部分看,首先是image_features.norm(dim=-1, keepdim=True),对特征向量的最后一维进行了归一化,但是保持相同的维度,此处保持形状不变便于接下来的计算。

接下来用text_features向量自身又除以了刚刚归一化了的向量,变为了单位向量,模长为1,便于评估。

接下来第八行将图像特征输入进提示学习模块,得到提示。

第十二行与第十三行做了跟视觉编码器部分相同的操作。

第十四行有点乱,拆开来看:

imf_i是图像特征,形状为(4,512)

为了与图像进行运算,需要对齐其形状,因此在这里做了转置。text_features.t()是对text_features向量作了转置,从(1000,512)转置为了(512,1000)

接下来放大因子作为计算权重,图像特征和转置后的文本特征作余弦相似度矩阵运算,得到相似度logits

那么到这里,CustomCLIP所关联的部分都看完了,可以回到最初的主类继续看看有什么遗漏的地方。

class CoCoOp(TrainerX)

CoCoOp主训练代码部分

还有一处定义模型的前向传播和反向传播部分:

    def forward_backward(self, batch):
        image, label = self.parse_batch_train(batch)

        model = self.model
        optim = self.optim
        scaler = self.scaler
        
        prec = self.cfg.TRAINER.COCOOP.PREC
        if prec == "amp":
            with autocast():
                loss = model(image, label)
            optim.zero_grad()
            scaler.scale(loss).backward()
            scaler.step(optim)
            scaler.update()
        else:
            loss = model(image, label)
            optim.zero_grad()
            loss.backward()
            optim.step()

        loss_summary = {"loss": loss.item()}

        if (self.batch_idx + 1) == self.num_batches:
            self.update_lr()

        return loss_summary

首先是将训练数据集进行了处理,提取出每个batch的图像和标签,另外定义了模型,优化器,和权重放大器,配置上,优化器是SGD,权重放大器没有定义。

接下来是精度的设定,因为配置是fp16而不是amp,所以其实只有用到下面比较标准的训练代码。

最后计算损失值,在每个epoch结束时更新学习率。

尾言

到此为止,对CoCoOp的主要组成部分都予以解析说明,算是完成了解读的工作,CoCoOp共计340行代码,以及关联的CLIP相关代码,初见时还是较为复杂,但是如果对CLIP甚有了解,那么CoCoOp的部分其实并没有太多,读起来也就没那么费劲了。

阅读完CoCoOp和CLIP的代码,再去阅读基于CLIP的代码就会轻松很多,后续我也会再写一篇阅读基于Dassl.Pytorch和CLIP的代码实现。

希望这篇文章对你的工作有所帮助。

Ref

  1. https://github.com/KaiyangZhou/Dassl.pytorch?tab=readme-ov-file#write-a-new-trainer
  2. https://medium.com/@howard022619/pytorch-register-buffer-%E7%AD%86%E8%A8%98-8f948573b735
  3. https://discuss.pytorch.org/t/nn-transformerencoderlayer-input-output-shape/99360
  4. https://huggingface.co/spaces/clip-italian/clip-italian-demo#logit-scale

No Comments

Send Comment Edit Comment


				
|´・ω・)ノ
ヾ(≧∇≦*)ゝ
(☆ω☆)
(╯‵□′)╯︵┴─┴
 ̄﹃ ̄
(/ω\)
∠( ᐛ 」∠)_
(๑•̀ㅁ•́ฅ)
→_→
୧(๑•̀⌄•́๑)૭
٩(ˊᗜˋ*)و
(ノ°ο°)ノ
(´இ皿இ`)
⌇●﹏●⌇
(ฅ´ω`ฅ)
(╯°A°)╯︵○○○
φ( ̄∇ ̄o)
ヾ(´・ ・`。)ノ"
( ง ᵒ̌皿ᵒ̌)ง⁼³₌₃
(ó﹏ò。)
Σ(っ °Д °;)っ
( ,,´・ω・)ノ"(´っω・`。)
╮(╯▽╰)╭
o(*////▽////*)q
>﹏<
( ๑´•ω•) "(ㆆᴗㆆ)
😂
😀
😅
😊
🙂
🙃
😌
😍
😘
😜
😝
😏
😒
🙄
😳
😡
😔
😫
😱
😭
💩
👻
🙌
🖕
👍
👫
👬
👭
🌚
🌝
🙈
💊
😶
🙏
🍦
🍉
😣
Source: github.com/k4yt3x/flowerhd
颜文字
Emoji
小恐龙
花!
Previous