论文地址：🐰

何凯明大神之作，通过无监督对比学习预训练Image Encoder的表征能力。后也被许多VLP算法作为ITC的底层算法来使用。

一方面由于源代码本身并不复杂，但是要求多GPU分布式训练，以及需要下载ImageNet这个大规模的数据集；另一方面 本次只是测试和阅读算法原理的实现，并不完整使用。因此，重写了一个低配版（流程不变，超参数没有严格要求设置，单GPU跑，数据集自己配置，几十张图片， no Shuffling BN）。

queue 即文中所构建的字典，起名为这个就是因为 C++ 中 的queue 容器，因为它是一种先进先出的数据结构。

目录

一、数据预处理

二、前向传播

网络结构

算法流程

---

一、数据预处理

对同一张图片进行数据增强操作，得到 query 和 key。

增强操作包括

transforms.RandomResizedCrop(224, scale=(0.2, 1.0)),transforms.RandomGrayscale(p=0.2),transforms.ColorJitter(0.4, 0.4, 0.4, 0.4),transforms.RandomHorizontalFlip(),normalize,

所以，dataloader中的每个输入样本是一个样本对儿。

通过下列方法实现

class TwoCropsTransform:"""Take two random crops of one image as the query and key."""def __init__(self, base_transform):self.base_transform = base_transformdef __call__(self, x):q = self.base_transform(x)k = self.base_transform(x)return [q, k]

二、前向传播

网络结构

代码中 encoder q 和 encoder k的网络结构用的都是ReNet 。ResNet最终的输出层包含了

(avgpool): AdaptiveAvgPool2d(output_size=(1, 1))(fc): Linear(in_features=2048, out_features=128, bias=True)

所以，输出的特征向量维度为 （N，C）。N为文中的Mini batch大小，代码中的超参数为batch size。C应该没有什么具体的含义，只是经验的设置为这一长度了（没找出来C的大小关乎什么）。

其输出还经过了L2归一化。 

算法流程

1、 q 送入 encoder q 得到输出，并经过L2归一化， (N,C)

2、 momentum 更新 key encoder。

3、 Shuffling BN（当然我重写的代码并没有实现这个，因为它需要多GPU，但这并不妨碍认识它的作用）

文中所述

大致意思由于ResNet使用了BN操作，因此由于Batch 数据之间的交互，使得模型利用它欺骗预设任务从而简单的找到一个低损失的解决方案，然而这个解决方案效果并不好，使得模型学习不到好的表征能力。

其提出的Shuffling BN

首先，把所有进程的Tensor的收集起来（如果分布式训练，一般每个GPU包含一个进程，所以收集的数据总量大小为 num GPUs * batch size），参考这里🤖

x_gather = concat_all_gather(x)

接下来制作打乱的索引，整个过程如下所示

    def _batch_shuffle_ddp(self, x):"""Batch shuffle, for making use of BatchNorm.*** Only support DistributedDataParallel (DDP) model. ***"""# gather from all gpusbatch_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 indexidx_shuffle = torch.randperm(batch_size_all).cuda()  # torch.randperm 将[0,n)数随机排列# broadcast to all gpustorch.distributed.broadcast(idx_shuffle, src=0)  # 将这个信息广播到所有其他进程# index for restoringidx_unshuffle = torch.argsort(idx_shuffle)  # 按照值大小顺序返回下标# shuffled index for this gpugpu_idx = torch.distributed.get_rank()  # 返回当前的进程idx_this = idx_shuffle.view(num_gpus, -1)[gpu_idx]  # idx_shuffle view 后 (num_gpus, batch size) 但是batch size中的索引是打乱顺序的return x_gather[idx_this], idx_unshuffle

最终返回 随机打乱顺序后挑选的当前进程的 batch size 大小的数据，也就是说进行 BN归一化后的数据已经不在 同一个原来的批 中了。

4、k 送入 encoder k 中，在经过L2 归一化， 和q一样。  （N，C）

5、Shuffling BN 对齐 q 和 k

如下面举例

# idx_shuffle
tensor([10, 16, 13,  2,  4,  0,  6, 21, 22, 31, 29,  3, 19, 17, 14, 30, 28, 12,24, 26,  8, 25, 11, 18,  5,  7, 27,  1, 15, 23, 20,  9])# idx_unshuffle
tensor([ 5, 27,  3, 11,  4, 24,  6, 25, 20, 31,  0, 22, 17,  2, 14, 28,  1, 13,23, 12, 30,  7,  8, 29, 18, 21, 19, 26, 16, 10, 15,  9])# q 的 idx_this
tensor([10, 16, 13,  2,  4,  0,  6, 21])# k 的 idx_this
tensor([ 5, 27,  3, 11,  4, 24,  6, 25])

这里主要关注的点是 这步是为了使 k对齐打乱顺序的q。q之前是打乱了顺序从而改变了每个batch的内容，相当于从所有的batch中随机挑选了 batch size的q，从而保证去除BN的影响。

而 k 不需要 再打乱了， 只需要从原有的batch size 数据分布中挑选出与q对应的数据即可。所以才在 shuffle BN q的过程中记录了indx unshuffle。

这里的对应关系举例，比如 index shuffle 中的 0 现在位于原来没打乱状态的索引 5处， 类似的 1 -->27, 2-->3, 以此类推。

注：不要被上面单进程的（即idx this）不对齐所迷惑，上面的只是分进程处理的，分布式训练最终会把所有进程的数据拼接起来一起处理，所以所有进程的数据对齐就行。

6、计算损失，即文中公式1

其中 用到的计算方法举例如下，分别用爱因斯坦求和公式实现，参考这里🤖

a = torch.tensor([[1, 2, 3], [1, 1, 1], [2, 2, 2]])
b = torch.tensor([[2, 2, 2], [2, 2, 2], [1, 1, 1]])
print(a)
print(b)
c = torch.einsum("nc, nc->n", [a, b])  # (3)
d = c.unsqueeze(-1)  # (3,1)
print(c)#=== 输出
tensor([[1, 2, 3],[1, 1, 1],[2, 2, 3]])
tensor([[2, 2, 2],[2, 2, 2],[1, 1, 1]])
tensor([12,  6,  7])
tensor([[12],[ 6],[ 7]])

a = torch.tensor([[1, 2, 3], [1, 1, 1], [2, 2, 3]])  # (3,3)
a1 = torch.tensor([[1, 2], [1, 1], [2, 2]])  # (3,2)
c = torch.einsum("nc,ck->nk", [a, a1])
print(a)
print(a1)
print(c)# ===输出
tensor([[1, 2, 3],[1, 1, 1],[2, 2, 3]])
tensor([[1, 2],[1, 1],[2, 2]])
tensor([[ 9, 10],[ 4,  5],[10, 12]])

这里的self.queue 即文中的字典 queue，初始化为

self.register_buffer("queue", torch.randn(dim, K))
self.queue = nn.functional.normalize(self.queue, dim=0)

K为字典的长度，默认设置65536。这里为什么设置为这个可能是由于ImageNet数据集比较大，所以设置的字典比较长，具体的长度设置好像没有做固定的要求，

来源于github官网。但代码中有要求，K必须是batch size 的倍数，这个为了确保字典的更新，方便执行入栈和弹出操作。这个字典像是C++的 queue容器的FIFO数据结构，即先进先出

self.K % batch_size == 0

        l_pos = torch.einsum("nc,nc->n", [q, k]).unsqueeze(-1)  #  (8,1)  对应元素相乘并第一维加和# negative logits: NxKl_neg = torch.einsum("nc,ck->nk", [q, self.queue.clone().detach()])  # (8,65536)  矩阵相乘# logits: Nx(1+K)logits = torch.cat([l_pos, l_neg], dim=1)  # (8,65537)# apply temperaturelogits /= self.Tlabels = torch.zeros(logits.shape[0], dtype=torch.long).cuda()  # (8,)loss = criterion(output, target)

这里看标签都是0，即第一个也就是0维数据为正样本。因为在拼接cat的时候正样本是在前面的。

7、更新字典

按mini batch 更新。具体地，如果 训练次数*mini batch size 小于字典长度，则字典queue每次都会填充新的key。若训练次数*mini batch size 大于 字典长度，则之前的被替换掉。

ptr = (ptr + batch_size) % self.K  # move pointer  8