对比学习(Contrastive Learning) (4)

《Unsupervised Feature Learning via Non-parametric Instance Discrimination》

2018年的cvpr论文，之后的许多对比学习的工作都有这篇文章的影子。

Introduction

作者通过观察监督学习的实验结果发现，视觉分类任务中，在softmax计算目标属于某一类的概率时，得分排在第二位的与得分排在第一位的类别在视觉上也会更接近，这种现象暗示出，一个好的学习方法能够认为相近的类别有大的相似度，如美洲豹与美洲虎；而不怎么有相似特征的类别，则相似度比较小，如美洲豹与书包。
进而，考虑到数据集中的每个样本实际上都是不一样的，在极端情况下，将每个样本视为一个类别，那么一个优秀的模型也能够按照对类别的相似度建模的方式来对不同样本之间的相似度建模，而如果真的按照样本个数来当做类别数目的话，对于ImageNet这种上百万张图片的数据集，计算softmax会非常的耗时，因此作者借鉴了noise-contrastive estimation （NCE）中的方法，近似的计算softmax，所以其实本质上，对比学习的损失函数就是想计算softmax，好多博客里面说对比学习的损失函数是在计算softmax，但是没有说为什么要这样设计。

Method

parametric classifier VS non-parametric classifier

假设有$n$张图片$x_1,x_2,...,x_n$，对它们进行特征提取之后的表示为$v_1,v_2,...,v_n$，softmax计算已知某个特征$v$属于第$i$类的条件概率为：
$$P(i|v)=\frac{\exp (w_i^{T}v)}{{\sum }_{j=1}^n\exp (w_j^{T}v)}$$
其中，$w_j^{T}v$反映了第$j$类与向量$v$的匹配程度，如果再进一步考虑，$w_j^T$反应$j$类的信息，如果$w_j^{T}v$的乘积越大，说明$w_j$与$v$的相似度越高，通过这种计算方式可以在空间上拉大不同类之间的距离，那么是否可以有一种不需要参数的表示形式，也能达到相近的效果呢？也就是说我们可不可以找到网络权重$w_j^T$的替代呢？
既然$w_j^T$反应$j$类的信息，那么有什么非参数的信息也同样可以表示第$i$类？所以作者用另一个变量$v_j$来替代权重，$v_j$表示第$j$个或第$j$类样本经过特征提取之后的向量，可以认为它跟$w_j^T$有相似的作用，这样做不仅减少了网络的训练参数量，已训练好的特征在下游任务上也会有更好的泛化能力，上式变为：
$$P(i|v)=\frac{\exp (v_i^{T}v/\tau )}{{\sum }_{j=1}^n\exp (v_j^{T}v/\tau )}$$
其中$\tau$用来控制分布的离散程度，它的作用不是我们探究的重点，不做进一步解释。根据监督学习中对于分类问题的损失函数计算：
$$J(\theta )=-\sum _{i=1}^n\log P(i|f_{\theta }(x_i))$$

如何选取不同的$v_j$?

对于整个数据集，初始化一个memory bank中存放每个样本经过特征提取之后的表示$V=v_j$，在训练过程中，随着网络逐渐收敛，不停地更新memory bank中的特征表示，这一部分直接看程序吧，语言表达能力有限 (⌐■_■)

上面的过程存在什么问题？

正如我们在introduction里面提到的，在计算$P(i|v)$的时候，分母上需要计算整个数据集中的每一个表征与当前样本表征的乘积，在计算上很耗时，所以作者借鉴了noise-contrastive estimation （NCE）和negative sampling来近似计算原始的softmax。
损失函数：
$$J_{NCE}(\theta )=-E_{P_d}[\log h(i,v)]-m{E}_{P_n}[\log (1-h(i,v^{\prime }))]$$
其中
$$h(i,v):=P(D=1|i,v)=\frac{P(i|v)}{P(i|v)+m{P}_n(i)}$$
$$P(i|v)=\frac{\exp (v^T{f}_i/\tau )}{Z_i}$$
$$Z_i=\sum _{j=1}^n\exp (v_j^T{f}_j/\tau )$$
其中$f_i$表示当前样本的特征，$v_j$表示从memory bank中随机选取的特征，$n$为均匀分布的噪声$P_n=1/n$，在实际情况下，不选取整个数据库中的样本做负样本，而是用蒙特卡洛来近似：
$$Z_i=n{E}_j[\exp (v_j^T{f}_i/\tau )]=\frac{n}{m}\sum _{k=1}^m\exp (v_{jk}^T{f}_i/\tau )$$

如何用训好的特征测试分类器？

假设有一张待测试的图像$\hat{x}$，计算其特征$\hat{f}=f_{\theta }(\hat{x})$，计算其与memory bank里面所有向量的相似度$s_i$，选择前$k$个相似度最高的样本${\mathcal{N}}_k$，计算这$k$个样本在每一个类别上的权重$w_c={\sum }_{i\in {\mathcal{N}}_k}{\alpha }_i\ast 1(c_i=c)$，其中${\alpha }_i=\exp (s_i/\tau )$，最后权重最高的那个类别为该测试图像的预测类别。

代码

官方的版本写的很啰嗦，这篇文章重点在损失函数的计算过程上，原版包括很多的实验验证，打开一个project的时候一上来找不到重点很难受，并且python和pytorch的维护环境做的不太友好，不同版本之间的pytorch在类上兼容性很差，近几年的深度学习都开始在分布式上面跑程序了，然而实验室还是每人一张卡，每次看到torch.distributed都自动略过 ⊙﹏⊙∥

预训练模型采用resnet-18，跟前几篇套路一样。重点看训练过程（包含损失函数的计算）和测试过程。

import argparse

import pandas as pd
import torch
import torch.nn.functional as F
import torch.optim as optim
from torch.optim.lr_scheduler import MultiStepLR
from torch.utils.data import DataLoader
from tqdm import tqdm

import utils
from model import Model

'''
    好像比官方版写的更像人话的样子
'''

# train for one epoch to learn unique features
def train(net, data_loader, train_optimizer):
    global z
    net.train()
    total_loss, total_num, train_bar = 0.0, 0, tqdm(data_loader)
    for data, target, pos_index in train_bar:
        data = data.to('cuda')
        train_optimizer.zero_grad()
        features = net(data)

        # randomly generate M+1 sample indexes for each batch ---> [B, M+1]
        # 随机生成0-50000的[btz, m + 1]个索引值，其中m为负样本的个数, m = 4096,算上正样本一共有4097个
        # [64, 4097]
        idx = torch.randint(high=n, size=(data.size(0), m + 1))
        # make the first sample as positive
        # 将idx索引为0位置处的索引值更改为该batch的数据在原数据集中的索引值，作为正样本，因为正样本一定是与其相似度最高的那个
        idx[:, 0] = pos_index
        # select memory vectors from memory bank ---> [B, 1+M, D]
        # memory_bank.shape = [50000, 128]
        # samples.shape = [64, 4097, 128]
        # 如何理解这里： 对于当前batch中的每一个样本经过encoder提取到的特征向量，与之相应地在memory_bank中随机选取4097个样本特征
        samples = torch.index_select(memory_bank, dim=0, index=idx.view(-1)).view(data.size(0), -1, feature_dim)

        # compute cos similarity between each feature vector and memory bank ---> [B, 1+M]
        # 两个三维矩阵的乘法
        # 这里这样理解：batch里面的每个样本都与4097个memory_bank中的样本求乘积，有一个作为正样本，另一个作为负样本
        # sim_matrix.shape = [64, 4097]
        sim_matrix = torch.bmm(samples.to(device=features.device), features.unsqueeze(dim=-1)).view(data.size(0), -1)
        # 在相似度矩阵上除以温度系数
        out = torch.exp(sim_matrix / temperature)
        # Monte Carlo approximation, use the approximation derived from initial batches as z
        # 原论文里面蒙特卡洛采样的部分 z = (n / m) \sum...
        # 这里存疑，均值应该是batch里面的每个样本对应的所有负样本求均值
        if z is None:
            z = torch.mean(out, dim=1) * n
            z = z.view(z.shape[0], -1).detach()
            # z = out.detach().mean() * n
        # compute P(i|v) ---> [B, 1+M]
        output = out / z
        # 上面在计算output的时候相当于每个元素都除了一个Z，也就是负样本与锚样本的乘积也除了一遍分母，但是实际计算P(i|v)时只需要计算锚样本与正样本
        # compute loss
        # compute log(h(i|v))=log(P(i|v)/(P(i|v)+M*P_n(i))) ---> [B]
        # 噪声为均匀分布
        p_d = (output.select(dim=-1, index=0) / (output.select(dim=-1, index=0) + m / n)).log()
        # compute log(1-h(i|v'))=log(1-P(i|v')/(P(i|v')+M*P_n(i))) ---> [B, M]
        p_n = ((m / n) / (output.narrow(dim=-1, start=1, length=m) + m / n)).log()
        # compute J_NCE(θ)=-E(P_d)-M*E(P_n)
        loss = - (p_d.sum() + p_n.sum()) / data.size(0)
        loss.backward(retain_graph=True)
        train_optimizer.step()

        # update memory bank ---> [B, D]
        # 更新memory_bank,这一个batch的数据中正样本的位置相应的由上一次的样本特征和这一次新生成的样本特征来决定，同时归一化保证绝对值为1
        pos_samples = samples.select(dim=1, index=0)
        pos_samples = features.detach().cpu() * momentum + pos_samples * (1.0 - momentum)
        pos_samples = F.normalize(pos_samples, dim=-1)
        memory_bank.index_copy_(dim=0, index=pos_index, source=pos_samples)

        total_num += data.size(0)
        total_loss += loss.item() * data.size(0)
        train_bar.set_description('Train Epoch: [{}/{}] Loss: {:.4f}'.format(epoch, epochs, total_loss / total_num))

    return total_loss / total_num


# test for one epoch, use weighted knn to find the most similar images' label to assign the test image
def test(net, memory_data_loader, test_data_loader):
    net.eval()
    total_top1, total_top5, total_num, feature_bank = 0.0, 0.0, 0, []
    with torch.no_grad():
        # generate feature bank
        # 构建一个feature bank，利用已经训练好的网络，将测试集的数据转化为特征
        for data, target, _ in tqdm(memory_data_loader, desc='Feature extracting'):
            feature_bank.append(net(data.to('cuda')))
        # [D, N]
        feature_bank = torch.cat(feature_bank).t().contiguous()
        # [N]
        feature_labels = torch.tensor(memory_data_loader.dataset.targets, device=feature_bank.device)
        # loop test data to predict the label by weighted knn search
        test_bar = tqdm(test_data_loader)
        for data, target, _ in test_bar:
            data, target = data.to('cuda'), target.to('cuda')
            output = net(data)

            total_num += data.size(0)
            # compute cos similarity between each feature vector and feature bank ---> [B, N]
            # 对于一个batch里面的所有样本，计算其与feature bank里面所有向量的相似度，选择结果最高的 k 个
            sim_matrix = torch.mm(output, feature_bank)
            # [B, K]
            sim_weight, sim_indices = sim_matrix.topk(k=k, dim=-1)
            # [B, K]
            sim_labels = torch.gather(feature_labels.expand(data.size(0), -1), dim=-1, index=sim_indices)
            sim_weight = (sim_weight / temperature).exp()

            # counts for each class
            one_hot_label = torch.zeros(data.size(0) * k, c, device=sim_labels.device)
            # [B*K, C]
            one_hot_label = one_hot_label.scatter(dim=-1, index=sim_labels.view(-1, 1), value=1.0)
            # weighted score ---> [B, C]
            pred_scores = torch.sum(one_hot_label.view(data.size(0), -1, c) * sim_weight.unsqueeze(dim=-1), dim=1)
            # 使用knn的方式选取最高的几个，如果预测索引和本身的索引相同，则视为预测正确
            pred_labels = pred_scores.argsort(dim=-1, descending=True)
            total_top1 += torch.sum((pred_labels[:, :1] == target.unsqueeze(dim=-1)).any(dim=-1).float()).item()
            total_top5 += torch.sum((pred_labels[:, :5] == target.unsqueeze(dim=-1)).any(dim=-1).float()).item()
            test_bar.set_description('Test Epoch: [{}/{}] Acc@1:{:.2f}% Acc@5:{:.2f}%'
                                     .format(epoch, epochs, total_top1 / total_num * 100, total_top5 / total_num * 100))

    return total_top1 / total_num * 100, total_top5 / total_num * 100


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Train NPID')
    parser.add_argument('--feature_dim', default=128, type=int, help='Feature dim for each image')
    parser.add_argument('--m', default=4096, type=int, help='Negative sample number')
    parser.add_argument('--temperature', default=0.1, type=float, help='Temperature used in softmax')
    parser.add_argument('--momentum', default=0.5, type=float, help='Momentum used for the update of memory bank')
    parser.add_argument('--k', default=200, type=int, help='Top k most similar images used to predict the label')
    parser.add_argument('--batch_size', default=64, type=int, help='Number of images in each mini-batch')
    parser.add_argument('--epochs', default=200, type=int, help='Number of sweeps over the dataset to train')

    # args parse
    args = parser.parse_args()
    feature_dim, m, temperature, momentum = args.feature_dim, args.m, args.temperature, args.momentum
    k, batch_size, epochs = args.k, args.batch_size, args.epochs

    # data prepare
    train_data = utils.CIFAR10Instance(root='../data', train=True, transform=utils.train_transform, download=True)
    train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True, drop_last=True)
    memory_data = utils.CIFAR10Instance(root='../data', train=True, transform=utils.test_transform, download=True)
    memory_loader = DataLoader(memory_data, batch_size=batch_size, shuffle=False, drop_last=True)
    test_data = utils.CIFAR10Instance(root='../data', train=False, transform=utils.test_transform, download=True)
    test_loader = DataLoader(test_data, batch_size=batch_size, shuffle=False, drop_last=True)

    # model setup and optimizer config
    model = Model(feature_dim).to('cuda')
    optimizer = optim.SGD(model.parameters(), lr=0.03, momentum=0.9, weight_decay=5e-4)
    print("# trainable model parameters:", sum(param.numel() if param.requires_grad else 0
                                               for param in model.parameters()))
    lr_scheduler = MultiStepLR(optimizer, milestones=[int(epochs * 0.6), int(epochs * 0.8)], gamma=0.1)

    # z as normalizer, init with None, c as num of train class, n as num of train data
    z, c, n = None, len(memory_data.classes), len(train_data)
    # init memory bank as unit random vector ---> [N, D]
    # 所有的样本构建一个memory_bank，初始化的时候bank里面是随机生成的向量，向量的维度与encoder的输出维度一致，||v|| = 1,所以需要标准化
    memory_bank = F.normalize(torch.randn(n, feature_dim), dim=-1)

    # training loop
    results = {'train_loss': [], 'test_acc@1': [], 'test_acc@5': []}
    best_acc = 0.0
    for epoch in range(1, epochs + 1):
        train_loss = train(model, train_loader, optimizer)
        results['train_loss'].append(train_loss)
        test_acc_1, test_acc_5 = test(model, memory_loader, test_loader)
        results['test_acc@1'].append(test_acc_1)
        results['test_acc@5'].append(test_acc_5)
        # save statistics
        data_frame = pd.DataFrame(data=results, index=range(1, epoch + 1))
        data_frame.to_csv('results/{}_results.csv'.format(feature_dim), index_label='epoch')
        lr_scheduler.step(epoch)
        if test_acc_1 > best_acc:
            best_acc = test_acc_1
            torch.save(model.state_dict(), 'epochs/{}_model.pth'.format(feature_dim))

数据库：

import torchvision.datasets as datasets
from PIL import Image
from torchvision import transforms


class CIFAR10Instance(datasets.CIFAR10):
    """CIFAR10Instance Dataset.
    """

    def __getitem__(self, index):
        img, target = self.data[index], self.targets[index]

        # doing this so that it is consistent with all other datasets
        # to return a PIL Image
        img = Image.fromarray(img)

        if self.transform is not None:
            img = self.transform(img)

        if self.target_transform is not None:
            target = self.target_transform(target)

        return img, target, index


train_transform = transforms.Compose([
    transforms.RandomResizedCrop(32, scale=(0.2, 1.0)),
    transforms.RandomApply([transforms.ColorJitter(0.4, 0.4, 0.4, 0.4)], p=0.8),
    transforms.RandomGrayscale(p=0.2),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))])

test_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))])