SimCLR/ResNet18小数批次机制失效：张量形状不兼容求助

问题现象 #

在7000张224×224 RGB图像的自定义数据集上训练SimCLR/ResNet18时：

• 批次大小设为100时训练正常（7000/100=70，无剩余批次）
• 批次大小设为32时，最后一批触发形状不兼容错误：

  RuntimeError: 张量a尺寸64与张量b尺寸48在非单例维度1不匹配

核心原因分析 #

错误中64=2×32、48=2×24，说明代码中硬编码了预设的批次大小（32），而非动态获取当前批次的实际样本数（最后一批为24）。SimCLR依赖每个样本生成两个增强视图，所有涉及批次维度的计算必须基于当前批次的真实大小，而非固定值。

排查与修复步骤 #

1. 检查数据加载环节 #

• 打印每个批次的两个增强视图张量形状，确认最后一批的样本数一致性：

  for x1, x2 in train_loader:
      print(f"x1 shape: {x1.shape}, x2 shape: {x2.shape}")
  

  若最后一批的x1.shape[0]与x2.shape[0]不一致，需修复Dataset的__getitem__方法，确保每个样本返回两个合法的增强视图。
• 确认DataLoader未设置drop_last=True（默认drop_last=False，保留最后一批），该参数仅能规避错误，无法解决根本问题。

2. 修复损失函数（NT-Xent）的硬编码问题 #

SimCLR的NT-Xent损失计算中，硬编码批次大小是最常见的出错点：

错误写法（硬编码固定值）

# 错误：使用预设batch_size而非实际批次大小
fixed_batch_size = 32
temperature = 0.5

z1 = encoder(x1)
z2 = encoder(x2)
z = torch.cat([z1, z2], dim=0)

# 硬编码生成64长度标签，与实际48长度的z不匹配
logits = torch.matmul(z, z.T) / temperature
labels = torch.arange(2 * fixed_batch_size, device=z.device)

正确写法（动态获取批次大小）

temperature = 0.5

z1 = encoder(x1)
z2 = encoder(x2)
current_batch_size = z1.size(0)
z = torch.cat([z1, z2], dim=0)

# 基于当前批次大小生成对应长度的标签
logits = torch.matmul(z, z.T) / temperature
labels = torch.arange(2 * current_batch_size, device=z.device)

3. 检查模型前向传播环节 #

确保ResNet18编码器的输出维度与输入批次大小匹配，避免自定义层中硬编码固定的批次维度值。

验证方案 #

修改代码后，单独提取最后一批数据运行训练逻辑，确认损失计算无形状错误后，再进行完整训练。

内容的提问来源于stack exchange，提问作者Willy Lutz