从零理解cmc对比学习算法：原理剖析与PyTorch实战

1次阅读

共计 3440 个字符，预计需要花费 9 分钟才能阅读完成。

在传统的监督学习中，模型需要大量标注数据才能训练出好的特征表示。但现实情况是，标注数据往往昂贵且耗时，这限制了模型的应用范围。自监督学习（Self-Supervised Learning）应运而生，它通过设计巧妙的预训练任务，让模型从未标注数据中学习有用的特征表示。

对比学习（Contrastive Learning）是自监督学习的一个重要分支，其核心思想是通过比较数据样本之间的相似性来学习特征表示。CMC（Contrastive Multiview Coding）算法是其中的一种代表性方法，它通过多视角特征对比显著提升了表征学习的效果。

CMC 算法的核心是 InfoNCE 损失函数，其数学表达式如下：

$$
\mathcal{L} = -\log \frac{\exp(sim(z_i, z_j)/\tau)}{\sum_{k=1}^{N} \exp(sim(z_i, z_k)/\tau)}
$$

其中，
– $z_i$ 和 $z_j$ 是同一图像的两个不同视角的特征表示（正样本对）
– $z_k$ 是其他图像的特征表示（负样本）
– $sim(\cdot,\cdot)$ 是相似度函数，通常使用余弦相似度
– $\tau$ 是温度系数，控制分布的尖锐程度
– $N$ 是负样本的数量

CMC 的关键创新在于利用多视角（Multi-View）特征对比。同一图像经过不同的数据增强（如裁剪、旋转、颜色变换等）会产生多个视角，CMC 强制这些视角的特征表示在嵌入空间中尽可能接近，同时远离其他图像的特征表示。

数据增强是对比学习成功的关键。我们需要为同一图像生成多个视角的变换版本。以下是典型的实现：

import torchvision.transforms as transforms

# 定义两种不同的数据增强方式
transform1 = transforms.Compose([transforms.RandomResizedCrop(size=224, scale=(0.2, 1.0)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomApply([transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)], p=0.8),
    transforms.RandomGrayscale(p=0.2),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

transform2 = transforms.Compose([transforms.RandomResizedCrop(size=224, scale=(0.2, 1.0)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomApply([transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)], p=0.8),
    transforms.RandomGrayscale(p=0.2),
    transforms.GaussianBlur(kernel_size=23, sigma=(0.1, 2.0)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

CMC 需要一个编码器网络（如 ResNet）和一个投影头（Projection Head）。投影头通常是一个简单的 MLP，用于将编码器输出的特征映射到对比学习空间。

import torch.nn as nn

class ProjectionHead(nn.Module):
    def __init__(self, input_dim=2048, hidden_dim=2048, output_dim=128):
        super().__init__()
        self.mlp = nn.Sequential(nn.Linear(input_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, output_dim)
        )

    def forward(self, x):
        return self.mlp(x)

下面是带温度系数的对比损失函数的 PyTorch 实现：

import torch
import torch.nn.functional as F

class ContrastiveLoss(nn.Module):
    def __init__(self, temperature=0.07):
        super().__init__()
        self.temperature = temperature

    def forward(self, features1, features2):
        # 归一化特征向量
        features1 = F.normalize(features1, dim=1)
        features2 = F.normalize(features2, dim=1)

        batch_size = features1.shape[0]
        # 计算相似度矩阵
        logits = torch.matmul(features1, features2.T) / self.temperature

        # 创建标签：对角线元素为正样本对
        labels = torch.arange(batch_size, device=features1.device)

        # 计算交叉熵损失
        loss1 = F.cross_entropy(logits, labels)
        loss2 = F.cross_entropy(logits.T, labels)

        return (loss1 + loss2) / 2

随着 batch size 增大，显存消耗会成为瓶颈。可以使用 memory bank 技术来存储历史特征表示，从而在不增加 batch size 的情况下获得更多的负样本。

class MemoryBank:
    def __init__(self, size, dim):
        self.size = size
        self.dim = dim
        self.bank = torch.randn(size, dim)
        self.ptr = 0

    def update(self, features):
        batch_size = features.shape[0]
        # 确保不会超出 memory bank 大小
        batch_size = min(batch_size, self.size - self.ptr)
        self.bank[self.ptr:self.ptr+batch_size] = features[:batch_size]
        self.ptr = (self.ptr + batch_size) % self.size

温度系数 $\tau$ 是一个关键超参数，控制着对比损失的 ” 硬度 ”：
– 较小的 $\tau$ 会使分布更尖锐，模型更关注困难的负样本
– 较大的 $\tau$ 会使分布更平滑，模型对所有负样本都 ” 一视同仁 ”

经验值通常在 0.05 到 0.2 之间。建议从小值开始尝试，观察损失曲线和下游任务性能。

在多 GPU 训练时，需要确保所有 GPU 上的特征表示能够同步计算相似度。PyTorch 的 DistributedDataParallel 可以帮助实现这一点。关键是要确保在计算对比损失前，收集所有 GPU 上的特征表示。

在 CIFAR-10 数据集上进行线性评估（Linear Evaluation）是验证对比学习效果的常用方法。具体步骤：
1. 使用 CMC 预训练一个特征提取器
2. 冻结特征提取器的权重，在上面训练一个线性分类器
3. 在测试集上评估分类准确率

典型结果（使用 ResNet-18 作为 backbone）：
– 随机初始化：约 50% 准确率
– 监督学习：约 95% 准确率
– CMC 预训练：可达到 85-90% 准确率

CMC 对比学习算法通过多视角特征对比，有效地从未标注数据中学习到了有意义的特征表示。本文详细介绍了其原理、PyTorch 实现以及实际应用中的优化技巧。虽然对比学习已经取得了显著进展，但在小数据集上的表现、计算效率等方面仍有改进空间。对于初学者来说，理解 CMC 算法是进入自监督学习领域的一个很好的起点。

正文完