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

1次阅读

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

在传统的监督学习中，模型训练严重依赖大量标注数据。然而，标注数据的获取成本高昂，特别是在医疗、遥感等领域。这使得监督学习在实际应用中受到限制。自监督学习（Self-Supervised Learning, SSL）应运而生，它通过设计巧妙的预训练任务，从无标注数据中学习有用的特征表示。对比学习（Contrastive Learning）是自监督学习的重要分支，其中 cmc（Contrastive Multiview Coding）算法通过多视角特征对比显著提升了表征学习效果。

cmc 算法的核心思想是最大化同一实例在不同视角下的特征相似性，同时最小化不同实例之间的特征相似性。其损失函数采用 InfoNCE（Noise Contrastive Estimation）形式，具体公式如下：

$$
\mathcal{L}{cmc} = -\log \frac{\exp(\text{sim}(z_i, z_j)/\tau)}{\sum
$$}^{N} \exp(\text{sim}(z_i, z_k)/\tau)

其中，$z_i$ 和 $z_j$ 是同一实例在不同视角下的特征表示，$\text{sim}(\cdot,\cdot)$ 是余弦相似度，$\tau$ 是温度系数，$N$ 是负样本数量。

以下是 cmc 算法的 PyTorch 实现代码，包含数据增强、投影头、对比损失计算等关键模块：

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms

# 数据增强模块
class Augmentation:
    def __init__(self):
        self.transform = 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])
        ])

    def __call__(self, x):
        return self.transform(x), self.transform(x)

# 投影头（MLP 实现）class ProjectionHead(nn.Module):
    def __init__(self, input_dim=2048, hidden_dim=2048, output_dim=128):
        super().__init__()
        self.net = 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.net(x)

# 带温度系数的对比损失计算
class ContrastiveLoss(nn.Module):
    def __init__(self, temperature=0.07):
        super().__init__()
        self.temperature = temperature

    def forward(self, z_i, z_j):
        # 计算相似度矩阵
        z = torch.cat([z_i, z_j], dim=0)
        sim_matrix = F.cosine_similarity(z.unsqueeze(1), z.unsqueeze(0), dim=2) / self.temperature

        # 构建正样本对和负样本对
        N = z_i.size(0)
        pos_mask = torch.zeros(2*N, 2*N, dtype=torch.bool, device=z_i.device)
        pos_mask[:N, N:] = True
        pos_mask[N:, :N] = True
        neg_mask = ~pos_mask

        # 计算损失
        exp_sim = torch.exp(sim_matrix)
        pos_sim = torch.sum(exp_sim * pos_mask, dim=1)
        neg_sim = torch.sum(exp_sim * neg_mask, dim=1)
        loss = -torch.log(pos_sim / neg_sim).mean()
        return loss

在对比学习中，负样本数量对模型性能至关重要。然而，随着负样本数量的增加，显存消耗也会急剧上升。以下是一些显存优化技巧：

Memory Bank：维护一个特征队列，存储历史批次的特征表示，作为额外的负样本来源。这样可以有效增加负样本数量，而不会显著增加显存消耗。
梯度检查点（Gradient Checkpointing）：在训练过程中，通过牺牲部分计算时间换取显存节省。PyTorch 中可以使用 torch.utils.checkpoint 实现。

温度系数 $\tau$ 是影响对比学习性能的关键超参数。以下是一些调参经验：