如何正确实现 InfoNCE 损失函数以避免张量形状不匹配问题

本文详解 infonce 损失函数中因标签构造依赖硬编码 batch_size 导致的 shape mismatch 错误，提供可扩展、设备无关的修复方案，并给出健壮的 pytorch 实现。

在自监督对比学习（如 SimCLR）中，InfoNCE 损失是核心组件，其正确性高度依赖于正负样本标签的精确构造。原始实现中常出现如下错误代码：

labels = torch.cat([torch.arange(self.args.batch_size) for i in range(self.args.n_views)], dim=0)

该写法隐式假设 features.shape[0] == self.args.n_views * self.args.batch_size，且 batch_size 始终能被整除——但当 batch_size=256、n_views=2 时，features 的实际 batch 维度应为 512，而若模型中间状态或数据加载存在微小偏差（如梯度累积、分布式训练中的局部 batch 不一致），self.args.batch_size 就不再可靠。此时 torch.arange(self.args.batch_size) 生成长度为 256 的序列，经拼接后仅得 512 元素，但后续 labels.unsqueeze(0) == labels.unsqueeze(1) 会生成 512×512 的布尔矩阵；而 mask = torch.eye(512) 也正确。真正出错点在于：当 features 张量因某些原因（如单卡调试、Dataloader drop_last=False 或混合精度缓存）实际 shape 为 [2, D] 时，labels.shape[0] 变为 2，mask 被创建为 2×2，但 labels[~mask] 尝试对 [2,2] 索引一个本应是 [512,512] 的张量，触发 PyTorch 的严格形状校验错误。

✅ 正确解法：完全抛弃对 self.args.batch_size 的依赖，直接从输入 features 的第一维动态推导真实批量大小。

由于每个样本生成 n_views 个增强视图，总特征数 features.shape[0] 必等于 n_views × actual_batch_size，因此每组样本的 ID 范围应为 torch.arange(features.shape[0] // self.args.n_views)：

def info_nce_loss(self, features):
    # ✅ 安全获取真实 batch size（不依赖 args）
    actual_batch_size = features.shape[0] // self.args.n_views
    labels = torch.cat(
        [torch.arange(actual_batch_size, device=features.device) 
         for _ in range(self.args.n_views)], 
        dim=0
    )

    # 构建对称标签矩阵：labels[i,j] == 1 当且仅当 i,j 属于同一原始样本的不同视图
    labels = (labels.unsqueeze(0) == labels.unsqueeze(1)).float()

    # L2 归一化确保余弦相似度计算
    features = F.normalize(features, dim=1)
    similarity_matrix = torch.matmul(features, features.T)

    # 创建对角掩码（排除自身与自身的相似度）
    mask = torch.eye(labels.shape[0], dtype=torch.bool, device=features.device)

    # 移除对角线元素，展平为 (N, N-1) 形状
    labels = labels[~mask].view(labels.shape[0], -1)
    similarity_matrix = similarity_matrix[~mask].view(similarity_matrix.shape[0], -1)

    # 提取正样本（lab
els==1）和负样本（labels==0）
    positives = similarity_matrix[labels.bool()].view(labels.shape[0], -1)
    negatives = similarity_matrix[~labels.bool()].view(similarity_matrix.shape[0], -1)

    # 拼接 logits：[pos_sim, neg_sim1, neg_sim2, ...]
    logits = torch.cat([positives, negatives], dim=1)
    # 标签恒为 0（所有正样本位于 logits 每行最左侧）
    labels = torch.zeros(logits.shape[0], dtype=torch.long, device=features.device)

    return logits / self.args.temperature, labels

⚠️ 关键注意事项：

• 务必指定 device：torch.arange(..., device=features.device) 避免 CPU/GPU 设备不匹配；
• 验证 n_views 整除性：在 __init__ 或 forward 开头添加断言 assert features.shape[0] % self.args.n_views == 0，防止静默错误；
• drop_last=True 推荐：DataLoader 中启用 drop_last=True，确保每个 batch 的样本数严格为 n_views × batch_size；
• 温度系数缩放：temperature 通常设为 0.1 或 0.2，过大会削弱对比强度，过小易导致梯度爆炸。

该实现彻底解耦了损失函数与训练配置参数，具备跨 batch size、多卡 DDP、梯度累积等场景的鲁棒性，是工业级对比学习代码库的标准实践。