同结构模型Keras训练正常但PyTorch始终预测全零问题排查

你的PyTorch代码存在4个和Keras实现不对齐的问题，其中核心问题会直接导致训练崩溃、输出全零：

• 核心错误：损失函数数值不稳定，直接引发ReLU神经元死亡
  Keras自带的binary_crossentropy内部做了两层数值稳定处理：先把Sigmoid输出裁剪到[1e-7, 1-1e-7]避免log(0)，再转换回logits用框架底层的稳定交叉熵实现计算，训练初期不会出现异常梯度。但PyTorch的F.binary_cross_entropy完全按照公式硬算，没有默认裁剪逻辑，一旦Sigmoid输出因为浮点精度落到0或1，就会产生inf/-inf的爆炸梯度，一次参数更新就会让大量ReLU神经元进入"死亡"状态（输入恒负、输出恒0、梯度恒0不再更新），最终模型输出全零。
• 梯度清零顺序不规范，存在梯度污染风险
  你当前的执行顺序是loss.backward() -> opt.step() -> opt.zero_grad()，首次迭代因为初始梯度为空可以正常运行，但如果训练中途出现梯度异常、或者迭代逻辑改动，残留的历史梯度会直接累加，导致参数更新完全错乱。标准实现需要把opt.zero_grad()放在每个batch训练步骤的最开头，先清空历史梯度再做前向计算。
• 数据类型未显式对齐，存在隐式转换风险
  Keras训练时会自动把输入特征、标签统一转换为和模型权重匹配的float32类型，但PyTorch的DataLoader不会做自动类型转换。如果你的输入是float64双精度类型、标签是int64长整型，轻则拖慢训练速度，重则在loss计算时触发类型提升、广播错误，间接导致训练崩溃。
• 初始化未完全对齐Keras逻辑
  PyTorch的nn.Linear默认bias使用均匀分布初始化，但Keras的Dense层bias默认初始化为全0。如果你只对权重做了Glorot Uniform初始化，没有手动把bias设为0，初始化分布的差异会让PyTorch模型在训练初期更容易产生大量负输入，提升ReLU死亡的概率。

修复后的对齐代码 #

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(19, 5),
            nn.ReLU(),
            nn.Linear(5, 5),
            nn.ReLU(),
            nn.Linear(5, 1)
            # 移除模型末尾的Sigmoid，改用数值稳定的带logits的BCE损失
        )
        # 完全对齐Keras初始化：权重Glorot Uniform，bias全0
        for layer in self.model:
            if isinstance(layer, nn.Linear):
                nn.init.xavier_uniform_(layer.weight)
                nn.init.zeros_(layer.bias)

    def forward(self, x):
        return self.model(x)


def train(train_ds, test_ds, model):
    train_dl = DataLoader(train_ds, batch_size=128, shuffle=True)
    test_dl = DataLoader(test_ds, batch_size=128, shuffle=True)
    opt = optim.Adam(model.parameters())
    # 用带logits的BCE，和Keras内部实现逻辑一致，数值稳定性更强
    loss_func = F.binary_cross_entropy_with_logits

    model.train()
    for epoch in range(30):
        for xb, yb in train_dl:
            # 显式转换数据类型为float32，和模型权重类型对齐
            xb = xb.float()
            yb = yb.float()
            # 每个batch先清空历史梯度
            opt.zero_grad()
            pred = model(xb)
            # 用squeeze代替reshape，更稳妥地去除多余维度对齐标签形状
            loss = loss_func(pred.squeeze(), yb)
            loss.backward()
            opt.step()

注意：预测阶段需要手动对模型输出做Sigmoid运算，才能得到0-1之间的概率值。

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