不妨試試MoCo，來替換ImageNet上pretrain模型！

無論是CV還是NLP領(lǐng)域，學(xué)習(xí)一個(gè)好的特征或者表達(dá)至關(guān)重要。對(duì)于NLP， 通過無監(jiān)督方法學(xué)習(xí)特征已經(jīng)取得好大的成功，但是在CV領(lǐng)域，目前主流的方案還是采用ImageNet上的有監(jiān)督pretrain模型。這是因?yàn)镹LP任務(wù)的輸入如words屬于一個(gè)離散空間，而CV任務(wù)的輸入圖片屬于一個(gè)高維連續(xù)空間。不過，目前的一些CV領(lǐng)域的無監(jiān)督學(xué)習(xí)方法也取得了一些較大的進(jìn)展，這里介紹的MoCo就是一種比較流行的無監(jiān)督學(xué)習(xí)方法。

對(duì)比學(xué)習(xí)

無監(jiān)督學(xué)習(xí)方法主要用來學(xué)習(xí)好的特征以用于特定的下游任務(wù)，這類方法往往需要建立一個(gè)代理任務(wù)（pretext task），所謂代理任務(wù)指的是這個(gè)任務(wù)本身并不是學(xué)習(xí)目的，真正想要的是通過這個(gè)任務(wù)學(xué)習(xí)到好的數(shù)據(jù)特征或者表達(dá)。這些代理任務(wù)總體可以分成兩大類：生成式（generative）和判別式（discriminative）。對(duì)于生成式方法，常見的通用自編碼器對(duì)圖像重建或者GAN，這類方法直接操作在pixel空間，計(jì)算費(fèi)時(shí)而且對(duì)于學(xué)習(xí)特征往往也不必要。而判別式方法和我們常用的監(jiān)督學(xué)習(xí)類似，通過特定的目標(biāo)函數(shù)來訓(xùn)練，但是輸入和標(biāo)簽都是來自于無標(biāo)注數(shù)據(jù)中，所以判別式方法也可以稱為自監(jiān)督學(xué)習(xí)（ self-supervised learning）。對(duì)比學(xué)習(xí)（contrastive learning）是判別式方法中最流行且成功的方法，其主要思路通過學(xué)習(xí)區(qū)分樣本對(duì)的特征來學(xué)習(xí)到好特征。

這里介紹的Moco方法就屬于對(duì)比學(xué)習(xí)的范疇，采用代理任務(wù)是instance discrimination task（見這篇paper），其主要點(diǎn)是最小化同一個(gè)圖片的不同view下的差異，這里的不同view是由同一個(gè)圖片進(jìn)行不同的數(shù)據(jù)增強(qiáng)得到。論文中將對(duì)比學(xué)習(xí)看成一個(gè)dictionary look-up task，這個(gè)dictionary的keys是一系列圖片經(jīng)過一個(gè)encoder得到的特征，這個(gè)encoder是一個(gè)網(wǎng)絡(luò)記為，對(duì)于輸入，其編碼后的特征為。另外，這里有另外一個(gè)encoder網(wǎng)絡(luò)，記為，它用來產(chǎn)生query的編碼，對(duì)于輸入，其編碼后得到（注意的是和可以是同一個(gè)網(wǎng)絡(luò)，也可以不同）。對(duì)比學(xué)習(xí)就是對(duì)比query和字典中的keys。這里假定對(duì)于一個(gè)特定的，這個(gè)字典中只有一個(gè)key與匹配，記為，又稱為positive key，在這里其實(shí)和是同一個(gè)圖片的不同view，而其他key是來自于不同的圖片，所以是negative key。那么對(duì)比學(xué)習(xí)的目標(biāo)函數(shù)就是要最大化和與之對(duì)應(yīng)的之間的相似度，而最小化和其他的之間的相似度。這里直接用兩個(gè)特征向量的點(diǎn)積來衡量相似度，而對(duì)比損失函數(shù)采用InfoNCE：

這里的是超參數(shù)，樣本中只有一個(gè)正例和K個(gè)負(fù)例，這個(gè)問題可以看成K+1分類問題，label是那個(gè)正例，所以損失函數(shù)就是基于softmax的交叉熵。

在實(shí)現(xiàn)上，對(duì)比學(xué)習(xí)共有兩種常用的實(shí)現(xiàn)機(jī)制：end-to-end和memory bank，如下圖所示。end-to-end方式應(yīng)該是最直觀的，無論是encoder 和encoder 均有梯度傳遞，雖然字典中keys一致性高（都是同參數(shù)的encoder產(chǎn)生），但是這種方式字典的大小一般是mini-batch的大小，因?yàn)槭芟抻贕PU顯存。memory bank是對(duì)end-to-end的改進(jìn)，memory bank包含數(shù)據(jù)集中所有樣本編碼后特征，這樣每個(gè)mini-batch的字典從memory bank中隨機(jī)采樣得到，這個(gè)過程不計(jì)算梯度，這樣字典的大小原則上可以很大，甚至和數(shù)據(jù)集大小一致。由于每個(gè)迭代，memory bank會(huì)更新mini-batch中的樣本，所以采樣得到的字典中keys一致性就差，因?yàn)榭赡軄碜圆煌?xùn)練step，此時(shí)encoder參數(shù)是不同的。

圖1 三種對(duì)比學(xué)習(xí)策略比較

Moco

無論是end-to-end還是memory bank都有優(yōu)缺點(diǎn)，一個(gè)最優(yōu)的方式是：字典的keys一致性更高，而且字典大小足夠大，這樣才有足夠多的負(fù)樣本以訓(xùn)練出好的encoder。Moco方法就是從這兩個(gè)方面來進(jìn)行改進(jìn)。Moco的核心是將字典看成一個(gè)隊(duì)列queue，隊(duì)列的大小可以比mini-batchh大，屬于超參數(shù)。隊(duì)列是逐步更新的在每次迭代時(shí)，當(dāng)前mini-batch的樣本入列，而隊(duì)列中最老的mini-batch樣本出列，這個(gè)字典中的keys始終是最新的。

使用隊(duì)列雖然可以使字典很大，但是與memory bank一樣，操作上難以通過BP來更新encoder，一個(gè)最簡單的方式直接用最新的作為，并且忽略梯度，但是實(shí)驗(yàn)中效果較差，這應(yīng)該和memory bank類似，字典中的keys的一致性較差，因?yàn)槊看蔚膃ncoder都更新很快。Moco采用momentum update策略來解決這個(gè)問題。假定和的參數(shù)分別是和，那么在訓(xùn)練的每次迭代中，通過BP來更新，但是采用momentum update：

這里的是momentum系數(shù)，一般取較大的值，論問中取0.999，這意味著更新過程非常平滑，這樣字典的keys一致性就更高。

Moco的核心就是維護(hù)一個(gè)動(dòng)態(tài)的queue作為字典，然后通過momentum update方式更新字典keys編碼網(wǎng)絡(luò)，其PyTorch偽代碼如下所示：

#?f_q,?f_k:?encoder?networks?for?query?and?key
#?queue:?dictionary?as?a?queue?of?K?keys?(CxK)
#?m:?momentum
#?t:?temperature

f_k.params?=?f_q.params?#?initialize
for?x?in?loader:?#?load?a?minibatch?x?with?N?samples
????x_q?=?aug(x)?#?a?randomly?augmented?version
????x_k?=?aug(x)?#?another?randomly?augmented?version
????
????q?=?f_q.forward(x_q)?#?queries:?NxC
????k?=?f_k.forward(x_k)?#?keys:?NxC
????k?=?k.detach()?#?no?gradient?to?keys
????
????#?positive?logits:?Nx1
????l_pos?=?bmm(q.view(N,1,C),?k.view(N,C,1))
????
????#?negative?logits:?NxK
????l_neg?=?mm(q.view(N,C),?queue.view(C,K))
????
????#?logits:?Nx(1+K)
????logits?=?cat([l_pos,?l_neg],?dim=1)
????
????#?contrastive?loss,?Eqn.(1)
????labels?=?zeros(N)?#?positives?are?the?0-th,?so?GT?label?is?0
????loss?=?CrossEntropyLoss(logits/t,?labels)
????
????#?SGD?update:?query?network
????loss.backward()
????update(f_q.params)
????
????#?momentum?update:?key?network
????f_k.params?=?m*f_k.params+(1-m)*f_q.params
????
????#?update?dictionary
????enqueue(queue,?k)?#?enqueue?the?current?minibatch
????dequeue(queue)?#?dequeue?the?earliest?minibatch
#?bmm:?batch?matrix?multiplication;?mm:?matrix?multiplication;?cat:?concatenation

對(duì)于同一張圖片，要通過不同的數(shù)據(jù)增強(qiáng)來產(chǎn)生和，Moco采用的數(shù)據(jù)增強(qiáng)方式如下所示：

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
]

論文中采用ResNet作為encoder，不過網(wǎng)絡(luò)的最后一個(gè)FC層輸出作為特征向量，其維度大小為128，這個(gè)向量經(jīng)過L2-norm進(jìn)行歸一化作為query和key。由于ResNet網(wǎng)絡(luò)存在BN層，但是直接采用BN層會(huì)惡化結(jié)果，因?yàn)锽N層中的mean和variance可能會(huì)泄露一些信息導(dǎo)致模型訓(xùn)練過程走捷徑，雖然loss很低，但是得到的特征卻并不好。Moco的解決方案是shuffling BN和多卡訓(xùn)練，具體操作如下：

#?compute?query?features，query計(jì)算過程不做改動(dòng)
q?=?self.encoder_q(im_q)??#?queries:?NxC
q?=?nn.functional.normalize(q,?dim=1)

#?key的計(jì)算要先shuffle?樣本，這樣同一張上query和key的BN層采用不同的樣本計(jì)算得到，避免信息泄露
with?torch.no_grads():
????#?shuffle?for?making?use?of?BN
????im_k,?idx_unshuffle?=?self._batch_shuffle_ddp(im_k)

????k?=?self.encoder_k(im_k)??#?keys:?NxC
????k?=?nn.functional.normalize(k,?dim=1)

????#?undo?shuffle
????k?=?self._batch_unshuffle_ddp(k,?idx_unshuffle)

對(duì)于這個(gè)問題，SimCLR采用的global BN策略，就是聚合所有卡上的BN層中的mean和variance，其它策略是將BN替換成layer Norm。

衡量無監(jiān)督方法得到的特征質(zhì)量常用的方法是Linear Classification Protocol，簡單來說，就是將encoder得到的特征凍結(jié)，后面加一個(gè)線性分類器來進(jìn)行監(jiān)督訓(xùn)練，在ImageNet數(shù)據(jù)集上，Moco比end-to-end和memory bank方法取得更好的效果。無監(jiān)督方法一個(gè)最重要應(yīng)用是將訓(xùn)練的特征遷移到其它下游任務(wù)，如檢測和分割，此時(shí)無監(jiān)督方法得到的encoder就作為初始化網(wǎng)絡(luò)來代替ImageNet數(shù)據(jù)集上監(jiān)督訓(xùn)練網(wǎng)絡(luò)，實(shí)驗(yàn)證明Moco確實(shí)可以取得更好的效果。

借鑒SimCLR方法的優(yōu)秀策略，Moco升級(jí)為MocoV2，首先采用更heavy的數(shù)據(jù)增強(qiáng)，數(shù)據(jù)增強(qiáng)在對(duì)比學(xué)習(xí)中至關(guān)重要，MocoV2采用的數(shù)據(jù)增強(qiáng)方式如下所示：

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
????????]

另外是將Moco中的fc head替換成一個(gè)2層的MLP head (隱含層為 2048-d，采用ReLU)，這個(gè)只是為了提升訓(xùn)練過程，linear classification和遷移時(shí)不會(huì)使用MLP。其它的改進(jìn)策略是采用cosine learning rate schedule。這些策略使得MocoV2效果進(jìn)一步提升：

參考

1. Momentum Contrast for Unsupervised Visual Representation Learning
2. Improved Baselines with Momentum Contrastive Learning
3. Unsupervised Feature Learning via Non-Parametric Instance-level Discrimination
4. A simple framework for contrastive learning of visual representations

   ---