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

前言

本文是参考大神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论文为例：

对比学习的具体做法为：

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相似。

如图所示：

其中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架构与之前对比学习类似，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_k$
   3.4. 将$x_q$ 和 $x_k$ 中元素“对应相乘”，得到128个数值。也就是得到了128个锚点与其对应的正样本的内积（相当于将向量内积作为相似度）
   3.5 将 $x_q$ 和 队列中的向量“两两相乘”，得到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$ 入队，响应的最早的一批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()