PyTorch实战：5步搞定监督对比学习（SupCon）损失函数实现-平芜编程栈

PyTorch实战：5步搞定监督对比学习（SupCon）损失函数实现

监督对比学习（Supervised Contrastive Learning）作为对比学习在监督场景下的扩展，正在计算机视觉、自然语言处理等领域展现出强大的特征提取能力。与传统的交叉熵损失相比，SupCon通过显式拉近同类样本、推远异类样本的特征表示，能够学习到更具判别性的嵌入空间。本文将聚焦PyTorch实现，用5个关键步骤带你从零实现SupCon损失函数。

1. 理解监督对比学习的核心思想

SupCon的核心创新在于巧妙利用了监督信息来定义正负样本。假设我们有一个batch中包含N个样本，每个样本经过两次不同的数据增强得到两个视图（views），那么：

正样本：与锚样本（anchor）类别相同的所有样本（包括不同视图）
负样本：与锚样本类别不同的所有样本

这种定义方式比自监督对比学习更直接地利用了标签信息。从数学上看，SupCon损失函数可以表示为：

$$ \mathcal{L}{sup} = \sum{i\in I}\frac{-1}{|P(i)|}\sum_{p\in P(i)}\log \frac{\exp(z_i \cdot z_p/\tau)}{\sum_{a\in A(i)}\exp(z_i \cdot z_a/\tau)} $$

其中：

$P(i)$是与样本$i$同类的正样本集合
$A(i)$是除$i$本身外的所有样本集合
$\tau$是温度系数，控制分布的尖锐程度

温度系数$\tau$的选择很关键：值太大会导致所有样本相似度趋同，太小则会使模型难以收敛。实践中通常设置在0.05到0.2之间。

2. 准备数据与特征编码器

在实现损失函数前，我们需要准备数据加载器和特征编码器。这里以CIFAR-10为例：

import torch import torchvision from torch import nn # 数据增强策略 train_transform = torchvision.transforms.Compose([ torchvision.transforms.RandomResizedCrop(32), torchvision.transforms.RandomHorizontalFlip(), torchvision.transforms.ColorJitter(0.4, 0.4, 0.4, 0.1), torchvision.transforms.ToTensor(), torchvision.transforms.Normalize( mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]) ]) # 加载CIFAR-10数据集 train_dataset = torchvision.datasets.CIFAR10( root='./data', train=True, transform=train_transform, download=True) # 简单的CNN编码器 class Encoder(nn.Module): def __init__(self): super().__init__() self.net = nn.Sequential( nn.Conv2d(3, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), nn.Flatten(), nn.Linear(64*8*8, 128) ) def forward(self, x): return self.net(x)

关键点在于数据增强策略的选择。SupCon的性能很大程度上依赖于使用的数据增强组合，常见的包括：

随机裁剪和大小调整
颜色抖动（亮度、对比度、饱和度、色调）
随机灰度化
高斯模糊

3. 实现SupCon损失函数

现在我们来逐步实现SupCon损失函数。完整的实现需要考虑以下几个技术细节：

相似度矩阵的高效计算
正负样本掩码（mask）的构建
数值稳定性处理
多视图支持

class SupConLoss(nn.Module): def __init__(self, temperature=0.07, contrast_mode='all'): super().__init__() self.temperature = temperature self.contrast_mode = contrast_mode def forward(self, features, labels=None): device = features.device # 特征维度处理 if len(features.shape) < 3: raise ValueError('特征需要是[bsz, n_views, ...]格式') batch_size = features.shape[0] # 构建标签掩码 labels = labels.view(-1, 1) mask = torch.eq(labels, labels.T).float().to(device) # 处理多视图特征 contrast_feature = torch.cat(torch.unbind(features, dim=1), dim=0) if self.contrast_mode == 'all': anchor_feature = contrast_feature anchor_count = features.shape[1] else: raise ValueError('不支持的对比模式') # 计算相似度矩阵 anchor_dot_contrast = torch.matmul( anchor_feature, contrast_feature.T) / self.temperature # 数值稳定性处理 logits_max, _ = torch.max(anchor_dot_contrast, dim=1, keepdim=True) logits = anchor_dot_contrast - logits_max.detach() # 构建排除自身的掩码 logits_mask = torch.scatter( torch.ones_like(mask), 1, torch.arange(batch_size * anchor_count).view(-1, 1).to(device), 0 ) mask = mask.repeat(anchor_count, anchor_count) * logits_mask # 计算log概率 exp_logits = torch.exp(logits) * logits_mask log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True)) # 计算正样本的平均log概率 mean_log_prob_pos = (mask * log_prob).sum(1) / mask.sum(1) # 最终损失 loss = -mean_log_prob_pos.mean() return loss

实现细节：使用torch.scatter构建的对角线掩码可以高效地排除样本与自身的对比，这是对比学习中常见的技巧。

4. 训练流程与技巧

有了损失函数后，我们需要设计完整的训练流程。以下是关键训练步骤：

前向传播：对每个样本生成两个增强视图
特征提取：通过编码器获取特征表示
损失计算：使用SupCon损失函数
反向传播：更新模型参数

def train_one_epoch(model, loss_fn, optimizer, loader, device): model.train() total_loss = 0 for images, labels in loader: images = torch.cat([images[0], images[1]], dim=0).to(device) labels = labels.to(device) # 获取特征 features = model(images) features = features.view(2, -1, features.size(-1)).permute(1, 0, 2) # 计算损失 loss = loss_fn(features, labels) # 反向传播 optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() return total_loss / len(loader)

训练SupCon模型时，有几个实用技巧值得注意：

学习率预热：前几个epoch使用较小的学习率，然后逐步增大
大batch size：对比学习通常需要较大的batch size（256以上）以获得足够的负样本
投影头：在编码器后添加一个小型MLP（如两层感知机）作为投影头，可以提升性能

5. 评估与应用

训练完成后，我们可以评估学习到的特征表示质量。常见评估方式包括：

线性评估协议：冻结特征提取器，只训练线性分类器
最近邻分类：在特征空间中使用k-NN分类
可视化分析：使用t-SNE或UMAP降维可视化特征分布

def evaluate(model, test_loader, device): model.eval() total_correct = 0 with torch.no_grad(): for images, labels in test_loader: images = images.to(device) labels = labels.to(device) features = model(images) preds = features.argmax(dim=1) total_correct += (preds == labels).sum().item() return total_correct / len(test_loader.dataset)

在实际项目中，SupCon学习到的特征可以用于：