二分类任务测试阶段CrossEntropyLoss损失值计算报错求助

一、问题根源分析 #

1. 数据格式不匹配问题 #

CrossEntropyLoss对输入有严格的格式要求，二分类场景下常见错误原因包括：

• 模型输出不符合要求：模型最后一层若输出单维度张量（如[batch_size, 1]，配合Sigmoid使用），但CrossEntropyLoss需要未经过Softmax的logits，形状必须为[batch_size, 2]（对应两个类别）。
• 标签格式错误：标签需为类别索引（0或1），形状为[batch_size]且类型为torch.long；若使用one-hot编码的标签（形状[batch_size, 2]），会直接触发维度不匹配错误。
• 设备/类型不统一：模型输出在GPU但标签在CPU，或标签为float类型而非long类型，也会导致格式校验失败。

2. RuntimeError: cuDNN error: CUDNN_STATUS_MAPPING_ERROR #

该错误属于GPU底层内存映射问题，常见诱因：

• GPU内存不足/碎片化：测试阶段batch size过大，或未关闭梯度计算导致内存占用过高，cuDNN无法完成张量的内存映射。
• 张量非连续存储：经过转置、切片等操作后的张量内存不连续，cuDNN优化的算子无法正常处理。
• 版本不兼容：PyTorch、CUDA、cuDNN版本不匹配，导致底层调用cuDNN接口时出现异常。
• 频繁设备切换：测试过程中数据在CPU和GPU之间频繁拷贝，引发内存映射混乱。

二、优化建议 #

1. 解决数据格式问题 #

• 规范输入格式：
  • 二分类任务中，模型最后一层不要添加Softmax，输出形状固定为[batch_size, 2]。
  • 将标签转换为类别索引：若原标签是one-hot格式，用torch.argmax(labels, dim=1).long()转换；确保标签形状为[batch_size]，类型为torch.long。
  • 统一设备：在计算损失前，将模型输出和标签同步到同一设备（GPU/CPU）：

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    logits = logits.to(device)
    labels = labels.to(device)
    

• 添加格式校验：在损失计算前加入断言，提前排查格式问题：

  batch_size = logits.shape[0]
  assert logits.shape == (batch_size, 2), f"Logits shape mismatch: expected ({batch_size},2), got {logits.shape}"
  assert labels.shape == (batch_size,), f"Labels shape mismatch: expected ({batch_size},), got {labels.shape}"
  assert labels.dtype == torch.long, f"Labels dtype mismatch: expected torch.long, got {labels.dtype}"
  

2. 解决cuDNN映射错误 #

• 优化内存使用：
  • 测试阶段强制关闭梯度计算，减少内存占用：

    with torch.no_grad():
        logits = model(test_loader)
        loss = criterion(logits, labels)
    

  • 适当减小测试batch size，避免内存溢出；若需保留大batch，可开启梯度检查点（但测试阶段不推荐）。
  • 定期清理GPU缓存：在测试循环的间隙（如每个epoch结束后）调用torch.cuda.empty_cache()，但不要频繁调用以免影响效率。
• 保证张量连续：对经过变形操作的张量调用.contiguous()，确保内存连续：

  logits = logits.contiguous()
  labels = labels.contiguous()
  

• 校验版本兼容性：对照PyTorch官方文档，确保PyTorch、CUDA、cuDNN版本匹配（例如PyTorch 2.1推荐搭配CUDA 11.8/12.1，cuDNN 8.7+）。
• 减少设备切换：修改数据加载器，直接将数据加载到GPU上，避免CPU与GPU之间的频繁拷贝：

  class TestDataset(Dataset):
      def __getitem__(self, idx):
          data, label = self.data[idx], self.labels[idx]
          return torch.tensor(data).to(device), torch.tensor(label, dtype=torch.long).to(device)
  

内容的提问来源于stack exchange，提问作者蘇煥淇