本文详解 infonce 损失实现中因标签生成逻辑硬编码 `batch_size` 导致的 shape mismatch 错误，指出根本原因在于 `labels` 构建未与实际特征维度对齐，并提供鲁棒、可扩展的修复方案。
在自监督对比学习（如 SimCLR）中，InfoNCE 损失的核心是构建正负样本对关系矩阵。常见错误源于标签（labels）张量的构造方式与输入特征的实际批量维度脱节。原始代码中：
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 严格成立，且 self.args.batch_size 始终准确反映单视图样本数。但当 batch_size=256、n_views=2 时，features 实际形状应为 [512, D]；若因数据加载、梯度累积或分布式训练导致 features.shape[0] 异常（例如仅含 2 个样本），torch.eye(labels.shape[0]) 就会生成 [2,2] 掩码，而 labels[~mask] 尝试索引一个本应为 [512, 512] 的展平矩阵——从而触发 “shape mismatch at index 0” 的 RuntimeError。
✅ 正确解法：始终从 features 的实际形状动态推导标签基数
应使用 features.shape[0] // self.args.n_views（而非 self.args.batch_size）作为每个视角的样本数，确保标签维度与特征一致：
def info_nce_loss(self, features):
# ✅ 动态计算每视角样本数，消除 batch_size 依赖
batch_per_view = features.shape[0] // self.args.n_views
labels = torch.cat(
[torch.arange(batch_per_view) for _ in range(self.args.n_views)],
dim=0
)
# 构建对称标签矩阵：label[i,j] == 1 当且仅当 i,j 为同一语义样本的不同增强视图
labels = (labels.unsqueeze(0) == labels.unsqueeze(1)).float()
labels = labels.to(self.args.device)
# L2 归一化特征，便于余弦相似度计算
features = F.normalize(features, dim=1)
similarity_matrix = torch.matmul(features, features.T) # [N, N], N = n_views * batch_per_view
# 创建对角掩码，移除自相似项（i==j）
mask = torch.eye(labels.shape[0], dtype=torch.bool).to(self.args.device)
# 应用掩码：展平后按行重塑，保持每行对应一个样本的其他所有样本
labels = labels[~mask].view(labels.shape[0], -1) # [N, N-1]
similarity_matrix = similarity_matrix[~mask].view(similarity_matrix.shape[0], -1) # [N, N-1]
# 提取正样本（同一语义的其他视图）和负样本（其余所有）
positives = similarity_matrix[labels.bool()].view(labels.shape[0], -1) # [N, 1] or [N, n_pos]
negatives = similarity_matrix[~labels.bool()].view(similarity_matrix.shape[0], -1) # [N, N-1-n_pos]
# 拼接 logits：每行形如 [pos_score, neg_score_1, neg_score_2, ...]
logits = torch.cat([positives, negatives], dim=1) # [N, 1 + (N-1-n_pos)]
labels = torch.zeros(logits.shape[0], dtype=torch.long).to(self.args.device) # 所有正样本位于第 0 列
return logits / self.args.temperature, labels
⚠️ 关键注意事项：
- 务必校验 features.shape[0] % self.args.n_views == 0，否则 batch_per_view 会截断。可在 forward 中添加断言：assert features.size(0) % self.args.n_views == 0。
- 若使用混合精度（AMP）或梯度检查点，确保 features 在归一化前已同步至正确设备与数据类型。
- self.args.n_views 必须与数据增强流水线输出的视图数严格一致（如 SimCLR 默认为 2）。
- 对于多卡 DDP 训练，features 已经是跨卡拼接结果，上述逻辑依然适用——因为 features.shape[0] 天然反映全局批量大小。
此修复方案彻底解耦了损失函数与配置参数的隐式绑定，提升代码鲁棒性与可移植性，适用于任意 batch_size、n_views 及分布式训练场景。
