使用PyTorch BCEWithLogitsLoss时遭遇维度不匹配错误求助

核心问题1：损失函数与模型输出逻辑冲突 #

你使用的nn.BCEWithLogitsLoss内部已集成Sigmoid运算，但你的模型最后又手动添加了self.sigmoid(x)处理输出，这不仅会导致计算逻辑冗余，也是触发维度广播错误的关键诱因之一。

核心问题2：权重维度不匹配 #

报错中的broadcast shape [32,2]表明，你传入损失函数的weights参数维度为[2]，但模型输出和标签的维度都是[32,1]，损失函数尝试将权重广播至[32,2]时，与输出维度[32,1]无法匹配，从而引发错误。

具体修复步骤 #

步骤1：对齐损失函数与模型输出 #

二选一执行即可：

• 选项A：保留模型的Sigmoid，替换损失函数为nn.BCELoss

  # 修改Trainer初始化中的损失函数定义
  self.cross_entropy = nn.BCELoss(weight=weights.to(self.device)).to(self.device)
  

• 选项B：移除模型中的Sigmoid，继续使用nn.BCEWithLogitsLoss

  # 修改模型forward方法，删除最后一行的Sigmoid调用
  def forward(self, x):
      x = self.fc1(x)
      x = self.relu(x)
      x = self.dropout(x)
      x = self.fc2(x)
      # 移除该行：x = self.sigmoid(x)
      return x
  

步骤2：调整权重维度适配单标签二分类 #

单标签二分类任务中，权重需与样本维度匹配：

• 若原weights为[w0, w1]（分别对应负、正样本权重），可提取正样本权重并调整维度：

  # 取正样本权重，调整为可广播的维度
  weights = weights[1].unsqueeze(0)
  

• 若使用nn.BCEWithLogitsLoss，更推荐用pos_weight参数单独指定正样本权重：

  # 初始化损失函数时指定pos_weight
  self.cross_entropy = nn.BCEWithLogitsLoss(pos_weight=weights[1].to(self.device)).to(self.device)
  

验证维度一致性 #

修复后需确认：

• preds.shape = torch.Size([32,1])
• labels.shape = torch.Size([32,1])
• 损失函数的权重/pos_weight维度与输入兼容

额外提示 #

• 单标签二分类场景下，模型最后一层输出维度设为1是合理的，对应二分类的概率值（经Sigmoid后）或logits（未经Sigmoid）。
• 训练时可打印weights.shape，快速排查维度匹配问题。

内容的提问来源于stack exchange，提问作者some nooby questions