【论文笔记】(对比学习经典论文MoCo) Momentum Contrast for Unsupervised Visual Representation Learning

文章目录

  • 前言
  • 1. 论文介绍
  • 2. 对比学习简介
  • 3. MoCo实现思路
  • 4. MoCo架构简介
  • 5. MoCo代码实现

前言

本文是参考大神bryanyzhu对MoCo论文的解读视频,并按照他的解读结合论文进行笔记总结。方便后续查阅。

论文地址 : https://arxiv.org/abs/1911.05722

论文代码(官方) : https://github.com/facebookresearch/moco

视频解读 : https://www.bilibili.com/video/av422340209

视频目录:

00:00 ~ 01:40 介绍作者
01:40 ~ 07:33 对比学习介绍
07:33 ~ 09:12 标题和作者介绍
09:12 ~ 12:03 摘要
12:03 ~ 13:23 导论: NLP和CV信号空间的区别。
13:23 ~ 17:08 导论: 将前人对比学习工作归纳成查询字典的任务
17:08 ~ 18:57 导论: 研究动机(字典大小和一致性问题)
18:58 ~ 22:05 导论:MoCo模型介绍
22:05 ~ 23:04 导论: MoCo训练任务的选择
23:04 ~ 25:30 导论:MoCo的结果
25:30 ~ 27:30 导论:讨论与结论
27:30 ~ 33:42 相关工作
33:42 ~ 42:25 方法:目标函数
42:45 ~ 49:50 详解队列(queue)和动量更新(Momentum update)
49:50 ~ 58:38 之前对比学习的局限性(① 端到端(end-to-end) 50:59 ~ 53:00 ② memory bank 53:00 ~ 56:33)
58:38 ~ 1:04:30 MoCo伪代码讲解
1:04:30 ~ 1:06:09 Shuffle BN
1:06:09 ~ 结尾 实验

本文假设你对对比学习有一定的了解,若不了解,可以参考这篇文章。

1. 论文介绍

本论文做的事情:提出了一种新的对比学习的方式,使用无监督学习来训练一个特征提取器

在之前,人们通常都是使用有监督学习任务的backbone作为特征提取器。例如,使用在ImageNet数据集上进行分类任务的backbone进行作为特征提取器进行迁移学习。而Moco可以利用无监督学习的方式训练一个特征提取器。这也就意味着:在CV领域,也可以实现像BERT一样,通过大量无标注的数据来训练一个优秀的特征提取器

MoCo的效果:拿MoCo训练出的特征提取器用于下游任务时,基本吊打拿有监督学习训练出的特征提取器。

2. 对比学习简介

本章简单介绍一下对比学习,为后续做铺垫。

对比学习的基本思路就是:让模型去区分哪些图片是类似的图片,哪些不是。最终希望的效果是,模型对相似图片提取的特征向量较为相近,不相似图片的特征向量距离较远。

以SimCLR论文为例:

【论文笔记】(对比学习经典论文MoCo) Momentum Contrast for Unsupervised Visual Representation Learning_第1张图片

对比学习的具体做法为:

  1. 取一个batch的图片,假设取16张图片。
  2. 对每一张图片进行如下操作:
    2.1 将该张图片进行一次随机数据增强,将其作为锚点。
    2.2 将该张图片再做一次随机数据增强,得到一个正样本。其与锚点为相似图片。
    2.3 将batch中的其他样本作为负样本。其与锚点为不相似图片。
  3. 经过第二步,我们会得到16+16x16张图片。其中每个锚点都有一张对样的正样本和15张对应的负样本。
  4. 我们将锚点通过一个Encoder提取特征
  5. 然后将正负样本通过另一个Encoder提取特征(也可以是同一个Encoder)
  6. 最后让锚点和其正负样本都计算相似度,预测哪一个是正样本。

3. MoCo实现思路

传统的对比学习可以归纳成一种任务:给定一个query和若干keys,找出这个query和以下那个key相似。

如图所示:

【论文笔记】(对比学习经典论文MoCo) Momentum Contrast for Unsupervised Visual Representation Learning_第2张图片

其中keys集合的大小就是Batch Size的大小。(Keys集合也称为“字典”)

如果想要训练出好的模型,那我们就需要如下两点:

  1. 字典越大越好
  2. 字典中的key的特征要保持一致性。换句话说:这些正负样本要保证是从一个Encoder中提取的特征。即使不是从一个Encoder提取的,也要是从相似的Encoder中提取的。

而之前的对比学习方法对于这两点没有办法做到兼容。例如对于上述提到的SimCLR的字典大小就是Batch Size大小,所以受限于显卡内存无法做到太大。 而一些其他论文虽然构建了一个大词典,但是词典中的特征向量都是来自不同阶段的Encoder,所以一致性不高。(不同阶段是指在训练过程中,Encoder参数是不断变化的)

与之前对比学习论文不同,在MoCo中,锚点称为query,而正负样本称为Key。

4. MoCo架构简介

MoCo主要就是解决了两个问题:

  1. 字典可以搞的很大
  2. 字典中的key的一致性很高。

MoCo的架构如下:

【论文笔记】(对比学习经典论文MoCo) Momentum Contrast for Unsupervised Visual Representation Learning_第3张图片

MoCo架构与之前对比学习类似,query和keys(就是锚点和正负样本)分别经过两个不同的Encoder,然后再计算相似度,最终预测字典中哪个key和query是一对儿。

MoCo与之前对比学习方法的区别主要在keys这边,主要有两点:

  1. Queue: MoCo维护了一个先进先出队列(FIFO queue),这个队列就是字典,里面存储的是图片的特征向量,所以可以很大。而队列的维护方式就是每次入队最新一批图片的特征向量,出队最早一批图片的特征向量。
  2. Momentum Encoder:为了保证queue中key的一致性,字典这边的encoder必须要缓慢的更新。所以作者采用了动量的方式。简单来说,就是Momentum Encoder = 0.99 * Momentum Encoder + 0.01 * Encoder。每次Encoder更新后,只拿它参数的0.01来更新Momentum Encoder,这样就可以保证Momentum Encoder更新缓慢了。

举个例子来描述整个过程:假设字典(队列)大小为65536, batch size为128, Encoder编码的特征向量大小为256.

  1. 初始化Encoder:首先初始化两个Encoder,一个作为Encoder,一个作为Momentum Encoder。然后让两个Encoder的参数保持一致。
  2. 初始化队列:拿65536张图片通过Momentum Encoder进行特征提取,然后放入队列。其中每个图片的特征为256.
  3. 接下来开始训练:
    3.1 从数据集中随机采样128张图片
    3.2 将这128张图片进行随机数据增强后送给Encoder,计算出其特征向量。
    3.3 将这128张图片再次随机数据增强后送给Momentum Encoder,计算其特征向量。此时我们就得到了两组特征向量,每组都有128个256维的特征向量。我们这里将其明明为 x q x_q xq x k x_k xk
    3.4. 将 x q x_q xq x k x_k xk 中元素“对应相乘”,得到128个数值。也就是得到了128个锚点与其对应的正样本的内积(相当于将向量内积作为相似度)
    3.5 将 x q x_q xq 和 队列中的向量“两两相乘”,得到128x65536个数值。也就是128个锚点和队列中所有负样本的内积。
    3.6 将锚点和与正样本的内积和与负样本的内积concat到一起,得到128x(1+65536)个内积。
    3.7 使用CrossEntropyLoss求出锚点与正样本对应的loss。即锚点与正样本越接近loss越小,锚点与负样本远离loss越大。
    3.8 反向传播更新encoder参数
    3.9 使用momentum机制更新Momentum Encoder,即Momentum Encoder = 0.99 * Momentum Encoder + 0.01 * Encoder
    3.10 将 x k x_k xk 入队,响应的最早的一批128个特征向量出队。
  4. 重复3的训练过程

5. MoCo代码实现

用论文中的伪代码表示整个过程为:

# 首先假设存在这么几个对象
# f_q, f_k: query和key的Encoder,也就是Encoder和Momentum Encoder
# queue: 字典队列。为一个Shape为CxK的Tensor。其中C为特征向量的维度,K为词典的大小。例如:256x65536
# m: momentum超参。例如 0.99
# t: temperature。温度超参。例如:0.05

f_k.params = f_q.params # 初始f_k的参数和f_q保持一致
for x in loader: # 从dataloader中拿出 N 个数据。x为一个batch的图片,例如x的shape为(128, 3, 224, 224)
	x_q = aug(x) # 对 x 进行随机的数据增强
	x_k = aug(x) # 对 x 进行另一种随机的数据增强
	
	q = f_q.forward(x_q) # 使用Encoder提取x_q的特征。q的shape为 NxC。 例如128x256, 即128张图片,每张图片的特征向量为256维
	k = f_k.forward(x_k) # 使用Momentum Encoder提取x_k的特征,q的shape也为 NxC
	k = k.detach() # Momentum Encoder不计算梯度
	
	# 求出锚点与正样本的内积。l_pos的shape为 Nx1。即每个锚点和对应正样本的向量内积(两两相乘)。例如 shape为128x1
	l_pos = bmm(q.view(N,1,C), k.view(N,C,1))

	# 让锚点和队列中所有的负样本计算内积。l_neg的shape为 NxK。例如 shape为128x65536
	l_neg = mm(q.view(N,C), queue.view(C,K))

	# 将正样本和负样本concat到一起,正样本放在最前面。logits的shape为 Nx(K+1)。例如 128x(65536+1)
	logits = cat([l_pos, l_neg], dim=1)

	labels = zeros(N) # 构造Label,因为正样本都是在0这个位置,所以构造N个0就行了。
	
	# 使用CrossEntropyLoss计算损失,这里需要除以温度
	loss = CrossEntropyLoss(logits/t, labels)
	
	# 反向传播并更新f_q的参数。
	loss.backward()
	update(f_q.params)

	# 使用momentum的方式更新f_k
	f_k.params = m*f_k.params+(1-m)*f_q.params
	
	# 更新字典
	enqueue(queue, k) # 将当前的 N 个样本的特征向量入队。
	dequeue(queue) # 同时再将最早的 N 个样本出队

论文中的伪代码还忽略了部分细节,这里我参考伪代码和论文源码再给出一个Pytorch可直接运行简单版,供大家参考:

import torch
from torch import nn
from copy import deepcopy

# 这里使用一个简单的卷积层+Dense层作为特征提取器
f_q = nn.Sequential(
    nn.Conv2d(3, 1, kernel_size=3, stride=5),
    nn.Flatten(),
    nn.Linear(2025, 256)
)
# 这里使用deepcopy的方式初始化Momentum Encoder,但实际应该使用参数赋值的方式
f_k = deepcopy(f_q)

m = 0.99  # 初始化Momentum超参
T = 0.05  # 初始化温度参数

# 初始化队列,队列大小为65536,队列中特征向量的维度为256
# 细节1:队列里的特征向量全都是单位向量,即向量的模为1
queue = torch.rand(256, 65536)

criteria = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(f_q.parameters(), lr=0.03)

# for x in datalaoder:
# 接下来模拟进行一次训练

# 初始化128张224x224的彩色图片
x = torch.rand(128, 3, 224, 224)
# 对图片进行两种不同的随机数据增强
x_q = x * torch.rand(128, 3, 224, 224)
x_k = x * torch.rand(128, 3, 224, 224)

# 使用Encoder提取x_q的特征
q = f_q(x_q)
# 细节2:对特征向量需要进行标准化,即将其变为单位向量。这样使用内积表示相似度时不会被向量长度所干扰。
q = nn.functional.normalize(q, dim=1)

# 接下来提取x_k的特征向量,由于Momentum Encoder不需要计算梯度,所以用no_grad修饰
with torch.no_grad():
    # 用Encoder的参数更新Momentum Encoder
    for param_q, param_k in zip(f_q.parameters(), f_k.parameters()):
        param_k.data = param_k.data * m + param_q.data * (1. - m)

    # 提取x_k的特征向量
    k = f_k(x_k)
    # 同样,需要对向量进行标准化,变为单位向量
    k = nn.functional.normalize(k, dim=1)

# q和k中的向量对应相乘。得到128个内积,即128个数字。l_pos: 128x1
# Einsum可以参考:https://blog.csdn.net/zhaohongfei_358/article/details/125273126
l_pos = torch.einsum('nc,nc->n', [q, k]).unsqueeze(-1)
# q和queue中的向量两两相乘。得到128x65536个内积
l_neg = torch.einsum('nc,ck->nk', [q, queue.clone().detach()])

# 将l_pos和l_neg拼起来。得到128x(1+65536)个内积
logits = torch.cat([l_pos, l_neg], dim=1)

# 将内积除以温度参数
logits /= T

# 构造Labels
labels = torch.zeros(logits.shape[0], dtype=torch.long)

# 将队列中的前128个数字出队,然后把k入队
queue = torch.concat([queue[:, 128:], k.permute(1, 0)], dim=1)

# 计算损失
loss = criteria(logits, labels)

# 反向传播,更新Encoder参数
loss.backward()
optimizer.step()
optimizer.zero_grad()

你可能感兴趣的:(机器学习,论文阅读,学习)