MoCo V1:视觉领域的自监督学习新突破

释放双眼,带上耳机,听听看~!
MoCo V1是视觉领域的自监督学习新突破,基于对比损失的动态字典库训练编码器,实现无监督预训练,为视觉领域的发展带来新的可能性。

自监督

监督学习在视觉领域应用广泛且成熟,但是也存在着一些挑战:监督学习需要大量带标签的数据,数据易得,但是给数据打标签工作枯燥且繁复,且成本高昂。我们希望能够省去人工标注的环节,使用无标签的数据集去预训练。

MoCo V1:视觉领域的自监督学习新突破

自监督学习是一种不需要人工输入数据标签的监督学习。不同于监督学习,自监督学习在没有人工输入的情况下,模型独自地分析数据,分类总结信息获得结果。相比于同样不需要标签的无监督学习,自监督学习不需要对数据进行分组和聚类。

MoCo V1:视觉领域的自监督学习新突破

自监督学习仅仅通过观察数据不同部分的相互作用来实现数据的学习表示,降低了对大量带标签数据的需求。在视觉领域,自监督学习包含两个任务,代理任务(pretext task)和下游任务(downstream task)。

  • 代理任务: 代理任务的目的是为了能够学习到能够应用到下游任务中的特征表示或模型权重。
  • 下游任务:下游任务可以是常见的,没有充足标签数据的视觉任务,比如,分类任务,检测任务和分割任务等等。

MoCo V1

前言

MoCo V1出现之前,无监督预训练在NLP领域大放异彩,比如GPT和BERT。但是在当时,在视觉领域,有监督预训练占据着绝对地位。究其原因,还是因为图像和语言文字信号空间的差异。MoCo之前视觉领域的无监督预训练基于对比损失(contrastive loss),建立动态的字典库(dynamic dictionaries)。字典中的”keys”是从数据中抽样,并被编码器编码。无监督学习训练编码器进行字典查找,一个被编码的“query”应当和它的匹配key相似,而与其他的“key”不同。

对比学习可以看做训练一个编码器进行字典查找。假设一个编码的query qq和字典中抽样出来的已编码的keys集合{k0,k1,k2,⋯ }{ k_{0}, k_{1}, k_{2}, cdots }。假如在字典中存在一个key 是qq匹配的, 记为k+k_{+}。对比损失是一个函数,当qq与其正key k+k_{+}相似且与所有其他key不同时,它的函数值是低的。将点乘值作为相似值,对比损失函数的一种形式,称为InfoNCE,如下所示。

Lq=−logexp(q⋅k+/τ)∑i=0Kexp(q⋅ki/τ)L_{q} = -log frac{expleft( q cdot k_{+} / tauright)}{sum_{i=0}^{K} expleft( q cdot k_{i} / tau right)}

其中τtau是温度参数,分母是一个正样本和KK个负样本的和。直观的看,这个损失是基于softmax的(K+1)left( K + 1 right)类分类器的log loss,将qq分类为k+k_{+}

通常,query的表示形式为q=fq(xq)q = f_{q} left( x^{q} right),其中fqf_{q}为编码网络,xqx^{q}是query样例。输入xqx^{q}xkx^{k}可以是图像,patches或者一组patches组成的上下文。fqf_{q}fkf_{k}网络可以相等,也可以部分共享,或者不同。

MoCo V1:视觉领域的自监督学习新突破

假设好的特征可以从包含丰富负样本集合的一个大字典中学习到,并且生成字典keys的编码器在训练变动过程中要保持一致性。上图展示三种keys如何维持和key 编码器如何更新的方式。第一种end-to-end方式,分别编码query和key的两个编码器通过反向传播进行更新。它使用当前mini-batch的样本作为字典。字典的大小与mini-batch的大小相关,受限于GPU内存大小。第二种是建立一个memory bank,memory bank包含数据集中的所有样本表示,每一个mini-batch编码key从memory bank中抽样,不更新,只有query的编码器通过反向传播进行更新。memory bank是能够支持大容量字典,然而,当最后一次看到样本时,它在memory bank 中的表示被更新,因此,抽样的keys基本上与以往的epoch的多个steps下的编码器有关,缺少一致性。MoCo基于memory bank,提出动量对比(Momentum Contrast),维护字典的大容量和一致性。

动量对比

MoCo V1:视觉领域的自监督学习新突破

Dictionary as a queue

动量对比的核心是将字典作为数据样本队列进行维护。队列的引入将字典大小与mini-batch的大小分离。字典的大小要远大于经典的mini-batch大小,并且可以灵活独立地将字典大小设为一个超参数。

字典中的样本逐渐被替代,当前mini-batch进入字典队列,在队列中的前mini-batch被移除字典队列。字典始终代表所有数据的抽样子集,而维护字典的额外计算是可控的。移除前mini-batch也是有益的,因为它的编码key已经过时了,与新的编码key缺少一致性。

# queue和queue_ptr的初始化
self.register_buffer("queue", torch.randn(dim, K))
self.queue = nn.functional.normalize(self.queue, dim=0)
self.register_buffer("queue_ptr", torch.zeros(1, dtype=torch.long))

# batch_size: 64, GPUs_num: 4
def _dequeue_and_enqueue(self, keys):
        # gather keys before updating queue
        # keys.shape  [16, 128]
        keys = concat_all_gather(keys) # 合并
        # keys.shape  [64, 128]
        batch_size = keys.shape[0]  # batch size: 64

        # ptr 队列当前指针头位置
        ptr = int(self.queue_ptr) 
        assert self.K % batch_size == 0  # for simplicity

        # replace the keys at ptr (dequeue and enqueue)
        self.queue[:, ptr:ptr + batch_size] = keys.T
        ptr = (ptr + batch_size) % self.K  # move pointer
        
        self.queue_ptr[0] = ptr

Momentum update

使用队列能保证字典的大容量,但是通过反向传播对key 编码器进行更新是非常棘手的。原始的解决方案是从query编码器中复制key 编码器fkf_{k},忽略它的梯度,实验表示这种做法效果很差。导致实验很差的原因,论文中猜想是由于过快地更新编码器减少了key 编码的一致性。为此,作者提出了动量更新解决这一问题。

fkf_{k}的参数表示为θktheta_{k}fqf_{q}表示为θqtheta_{q},以下述公式更新θktheta_{k}

θk←mθk+(1−m)θqtheta_{k} leftarrow m theta_{k} + left( 1 – m right) theta_{q}

其中m∈[0,1)]m in [ 0,1 )]是一个动量系数。只有参数θqtheta_{q}是通过反向传播进行更新。上述的动量公式使得θktheta_{k}更加平滑地改变。尽管在队列中的keys被不同mini-batches时刻的编码器进行编码,这些编码之间的差异也是小的。

@torch.no_grad()
    def _momentum_update_key_encoder(self):
        """
        Momentum update of the key encoder
        """
        for param_q, param_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):
            param_k.data = param_k.data * self.m + param_q.data * (1. - self.m)

代理任务

代理任务有各种各样,MoCo更关注地是损失函数部分。Moco选择了简单的实例辨别任务(instance discrimination task)作为代理任务。

如果一个query和一个key来自于相同的照片,则把它们单做是正样本,否则,看做负样本对。在随机数据增强下,对同一个图像随机进行两次查看,构建正样本对。queries和keys分别被它们的编码器进行编码,记为fqf_{q}fkf_{k},编码器可以是任意卷积网络。

MoCo V1:视觉领域的自监督学习新突破

算法1展示了代理任务的伪代码,对于当前的mini-batch,编码queries和它们对应的keys,作为正样本对,负样本位于队列中。

Shuffling BN

在实验中,作者发现BN能够阻止模型学习更好的特征,论文使用shuffling BN解决这个问题。

使用多块GPUs进行训练,在每块GPU上,单独地执行BN操作。对于key编码器fkf_{k},在分布给GPUs之前,随机打乱当前mini-batch的样本顺序。query编码器fqf_{q}中的样本顺序不变动。这样能够确保计算query和它的正样本的批处理统计数据来自不同的子样本集。

   @torch.no_grad()
    def _batch_shuffle_ddp(self, x):
        """
        Batch shuffle, for making use of BatchNorm.
        *** Only support DistributedDataParallel (DDP) model. ***
        """
        # gather from all gpus
        batch_size_this = x.shape[0]
        x_gather = concat_all_gather(x)
        batch_size_all = x_gather.shape[0]

        num_gpus = batch_size_all // batch_size_this

        # random shuffle index  随机打乱index或者序列
        idx_shuffle = torch.randperm(batch_size_all).cuda()
        # broadcast to all gpus  
        torch.distributed.broadcast(idx_shuffle, src=0) # src源进程
        # index for restoring  # 重新排序
        idx_unshuffle = torch.argsort(idx_shuffle)

        # shuffled index for this gpu  
        gpu_idx = torch.distributed.get_rank() # 返回当前进程的排名
        idx_this = idx_shuffle.view(num_gpus, -1)[gpu_idx]

        return x_gather[idx_this], idx_unshuffle

实验

论文中采用了ResNet作为编码器,ResNet在全局平均池化之后的最后一个全连接层有一个固定维度的输出,128−D128-D。输出然后进行L2L2正则化。正则化的输出才是query和key的特征表示。

数据增强设定为:从resized的图像随机裁减224×224224 times 224,然后再进行random color jittering, random horizontal flip 和 random grayscale conversion。

优化器采用SGD。

线性分类器:无监督预训练完成之后,冻结所有的特征,然后训练一个有监督的线性分类器,全连接层后接一个softmax。

MoCo V2

MoCo v2借鉴了SimCLR中的两个有效设计,MLP projection head 和 更多的数据增广。

MLP-head: 将MoCo中的fc head替换成2-layer MLP head。head仅在训练阶段起作用,线性分类器或者迁移阶段,不使用MLP。

数据增广:在原有的数据增广基础上增加了模糊增广方式。

MoCo V1:视觉领域的自监督学习新突破


# MoCo V1和MoCo v2的数据增广方式
if args.aug_plus:
        # MoCo v2's aug: similar to SimCLR https://arxiv.org/abs/2002.05709
        augmentation = [
            transforms.RandomResizedCrop(224, scale=(0.2, 1.)),
            transforms.RandomApply([
                transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)  # not strengthened
            ], p=0.8),
            transforms.RandomGrayscale(p=0.2),
            transforms.RandomApply([moco.loader.GaussianBlur([.1, 2.])], p=0.5),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize
        ]
    else:
        # MoCo v1's aug: the same as InstDisc https://arxiv.org/abs/1805.01978
        augmentation = [
            transforms.RandomResizedCrop(224, scale=(0.2, 1.)),
            transforms.RandomGrayscale(p=0.2),
            transforms.ColorJitter(0.4, 0.4, 0.4, 0.4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize
        ]

MoCo V3

MoCo V3探讨如何在ViT上训练MoCo。如果batch足够大,比如4096,memory queue所带来的的增益就会减小,MoCo v3舍弃了memory queue,这意味着MoCo v3中的负样本存在于同一个batch中。不同于MoCo v1/v2版本, MoCo v3 对于同一个图像的2个增强版本x1x_{1}x2x_{2},分别通过fqf_qfkf_{k}得到q1,q2q_{1}, q_{2}k1,k2k_{1}, k_{2},采用对称loss的策略,计算q1,k2q_{1}, k_{2}q2,k1q_{2}, k_{1}的对比loss和优化更新fqf_q的参数。

编码器fqf_q包含一个backbone、一个projection head和一个额外的prediction head。编码器fkf_{k}包含一个backbone和一个projection head。fkf_{k}的更新还和MoCo一样,进行动量更新。MoCo v3的伪代码如下所示:

MoCo V1:视觉领域的自监督学习新突破

def forward(self, x1, x2, m):
        """
        Input:
            x1: first views of images
            x2: second views of images
            m: moco momentum
        Output:
            loss
        """

        # compute features
        q1 = self.predictor(self.base_encoder(x1))
        q2 = self.predictor(self.base_encoder(x2))

        with torch.no_grad():  # no gradient
            self._update_momentum_encoder(m)  # update the momentum encoder

            # compute momentum features as targets
            k1 = self.momentum_encoder(x1)
            k2 = self.momentum_encoder(x2)

        return self.contrastive_loss(q1, k2) + self.contrastive_loss(q2, k1)
def contrastive_loss(self, q, k):
    # normalize
    q = nn.functional.normalize(q, dim=1)
    k = nn.functional.normalize(k, dim=1)
    # gather all targets
    k = concat_all_gather(k)
    # Einstein sum is more intuitive  爱因斯坦求和约定
    logits = torch.einsum('nc,mc->nm', [q, k]) / self.T  # self.T softmax temperature
    N = logits.shape[0]  # batch size per GPU
    labels = (torch.arange(N, dtype=torch.long) + N * torch.distributed.get_rank()).cuda()
    return nn.CrossEntropyLoss()(logits, labels) * (2 * self.T)

论文发现不稳定性是影响自我监督ViT训练的主要问题。但是不稳定ViT训练不会导致灾难性地失败,但是会稍微降低1∼3%1sim 3%的准确率。

MoCo V1:视觉领域的自监督学习新突破

作者通过实验探究了批次大小,学习率和优化器对模型稳定性的影响。从图像1中可以看出,批次大小从1k到2k,训练曲线平滑上升;当批次大小为4k时,训练曲线出现不稳定性;当批次大小为6k时,训练曲线出现大的下降。从图像2中可以看出,当学习率较小时,训练较稳定,但是欠拟合,随着学习率的升高,曲线组件不稳定,会出现较大的dip。使用AdamW优化器时,随着学习率的增加,曲线虽然平滑,但是出现了衰退。

MoCo V1:视觉领域的自监督学习新突破

如图4所示,在训练过程中,梯度的一个突然改变会导致训练曲线中的一个下降。比较了所有层的剃度峰值之后,观察到第一层(patch projection)会更早出现梯度峰值,然后最层蔓延到最后一层。作者猜测不稳定现象在浅层会更早产生,于是将patch projection冻住进行训练,也就是使用一个固定的random patch projection层嵌入patches,训练曲线变得平滑并且准确率提升。将此方法用到SimCLR和BYOL方法中,也出现了同样的效果。

参考

  1. Moco Code
  2. self-Supervised Learning
  3. Improved Baselines with Momentum Contrastive Learning
  4. MoCo V3 github
  5. An Empirical Study of Training Self-Supervised Vision Transformers
  6. Self-Supervised Learning – Pretext Tasks
本网站的内容主要来自互联网上的各种资源,仅供参考和信息分享之用,不代表本网站拥有相关版权或知识产权。如您认为内容侵犯您的权益,请联系我们,我们将尽快采取行动,包括删除或更正。
AI教程

Pytorch实战:波士顿房价预测模型搭建

2023-12-17 18:08:14

AI教程

MetaApp-推荐广告研发部:基于DeepRec的稀疏模型训练实践

2023-12-17 18:25:14

个人中心
购物车
优惠劵
今日签到
有新私信 私信列表
搜索