卷积神经网络输出指数级过大及TernausNet训练梯度消失等问题求助

Alright, let's tackle your questions one by one—first the CNN output explosion, then your PyTorch Lightning training headaches.

There are several common reasons behind this behavior:

• Missing normalization layers: Convolutional layers apply linear transformations, and without BatchNorm, LayerNorm, or similar normalization, the distribution of each layer's output can drift further and further from a stable range. Over time, this causes values to balloon exponentially, especially if your weights are initialized to larger values.
• Incorrect activation function usage: If you're using an activation that doesn't constrain output (like a linear layer with no activation) or if you're using an activation that saturates but inputs get so large that the output hits the saturation ceiling (e.g., Sigmoid approaching 1 for huge positive inputs), the network may keep amplifying outputs to minimize loss.
• Poor weight initialization: If you initialize weights with too high a variance (e.g., not using He/Xavier initialization for your activation type), each layer's output variance will grow exponentially as data passes through the network, leading to massive output values.
• Gradient explosion leading to runaway weights: During training, if gradients explode, weight updates become extremely large. This makes the next forward pass produce even bigger outputs, creating a vicious cycle.
• Mismatched loss function and output: If your loss function doesn't impose constraints on the output range (e.g., using a loss designed for probabilities but not applying a Sigmoid/Softmax), the network may learn to crank up outputs to extreme values to minimize loss.

Let's break down your core problems: zero gradients, outputs reaching 10^8, and the weird negative loss when switching to BCEWithLogitsLoss.

Why you're seeing zero gradients + massive outputs #

Looking at your DiceLoss code, the biggest red flag is double application of Sigmoid:
You mentioned your model uses a Sigmoid activation, but your loss function also runs torch.sigmoid(logits) on the input. Here's why that breaks everything:

• When your model's Sigmoid outputs are already near 1 (because logits are huge), applying Sigmoid again leaves the value almost unchanged at 1. The derivative of Sigmoid at these extreme values is almost zero—so when backpropagating, the gradient through the double Sigmoid becomes negligible, leading to zero gradients for all weights.
• To minimize the DiceLoss (which is 1 - dice), the network will try to make proba as close to targets as possible. If targets are mostly 1, the model will push logits to infinity to make Sigmoid hit 1—hence your outputs reaching 10^8.

Other contributing factors:

• Your eps is set to 0: When summation is small, dividing by near-zero can cause numerical instability, pushing the network to amplify outputs to avoid this.
• Lack of normalization: If TernausNet doesn't include BatchNorm layers, output distributions can drift, making values explode.

Fixes for gradient vanishing & output explosion #

1. Remove the duplicate Sigmoid in DiceLoss
   Since your model already applies Sigmoid, change your loss code to use the model's output directly:

   class DiceLoss(nn.Module):
       def __init__(self):
           super().__init__()
       def forward(self, proba, targets, eps=1e-6, threshold=None):
           # Remove this line: proba = torch.sigmoid(logits)
           proba = proba.view(proba.shape[0], 1, -1)
           targets = targets.view(targets.shape[0], 1, -1)
           if threshold:
               proba = (proba > threshold).float()
           intersection = torch.sum(proba * targets, dim=1)
           summation = torch.sum(proba, dim=1) + torch.sum(targets, dim=1)
           dice = (2.0 * intersection + eps) / (summation + eps)
           return (1 - dice).mean()
   

   Note: I renamed the input from logits to proba to reflect that it's already a probability value from your model's Sigmoid.

2. Add numerical stability with eps
   Set eps=1e-6 (not 0) to avoid division by zero and stabilize loss calculations.

3. Use gradient clipping
   In PyTorch Lightning, enable gradient clipping in your Trainer to prevent gradient explosions:

   Trainer(gradient_clip_val=0.5, ...)
   

4. Check weight initialization
   Ensure TernausNet uses appropriate initialization (like He initialization for ReLU layers) to prevent initial output values from being too large.

5. Add normalization layers
   If your network doesn't have BatchNorm2d after convolutional layers, add them to stabilize layer outputs and prevent value drift.

Why BCEWithLogitsLoss gives negative loss & won't converge #

BCEWithLogitsLoss is designed to work directly with logits (no Sigmoid in the model), and its output should always be non-negative. Negative loss means something's misaligned:

• You're still applying Sigmoid in the model: BCEWithLogitsLoss includes a Sigmoid internally. If your model already outputs probabilities, feeding them into BCEWithLogitsLoss will cause numerical chaos—leading to negative or infinite loss values.
• Incorrect target format: BCEWithLogitsLoss expects targets to be float tensors with values in [0, 1], matching the shape of your logits. If your targets are integers (e.g., LongTensor) or have values outside [0,1], loss calculations will break.
• Numerical overflow: If your logits are already at 10^8, sigmoid(logits) is effectively 1, and log(1 - sigmoid(logits)) becomes log(0) (negative infinity). The loss calculation can overflow and produce nonsensical negative values.

Fixes for BCEWithLogitsLoss issues #

1. Remove Sigmoid from your model
   Let BCEWithLogitsLoss handle the Sigmoid internally—this is more numerically stable anyway.

2. Fix your target tensors
   Convert targets to float32 and ensure they're in [0,1]. For example:

   targets = targets.float()  # if they were integers
   targets = targets / targets.max()  # if they're scaled incorrectly
   

3. Resolve the output explosion first
   Apply the fixes above (gradient clipping, normalization, correct initialization) to get logits back to a reasonable range before switching to BCEWithLogitsLoss.

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