灰度图像卷积模型输出动态范围压缩问题求助

核心问题诊断 #

1. 最后一层无输出约束：模型最后一层是纯卷积层，未添加激活函数约束输出范围，结合ReLU的特性，导致输出被挤压到极窄区间。
2. ReLU神经元死亡：ReLU在输入小于0时梯度为0，部分神经元会彻底停止更新，对应你观察到的“部分参数或层停止学习”现象。
3. 卷积层padding非对称：第一个卷积层kernel_size=13搭配padding=5会导致特征图尺寸非对称缩小，虽然后续层补回了尺寸，但可能引入不必要的信息丢失。
4. 训练循环变量错误：函数参数传入的是dataloader，但循环内硬编码使用train_dataloader，会导致复用全局变量，逻辑易出错。

具体修复方案 #

1. 调整模型结构 #

• 最后一层添加Sigmoid激活，约束输出到[0,1]区间，匹配输入与GT的归一化范围。
• 将ReLU替换为LeakyReLU，保留负区间梯度，避免神经元死亡。
• 修正卷积层padding为对称值，保证每一步特征图尺寸与输入完全一致：

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
           # padding=(kernel_size-1)//2，保证输出尺寸与输入一致
           nn.Conv2d(in_channels=1, out_channels=64, kernel_size=13, stride=1, padding=6),
           nn.LeakyReLU(0.1, inplace=True),
           nn.Conv2d(in_channels=64, out_channels=32, kernel_size=3, stride=1, padding=1),
           nn.LeakyReLU(0.1, inplace=True),
           nn.Conv2d(in_channels=32, out_channels=1, kernel_size=5, stride=1, padding=2),
           nn.Sigmoid()  # 约束输出到[0,1]区间
          )

    def forward(self, input_image):
        output_image = self.model(input_image)
        return output_image

2. 修复训练循环逻辑 #

将循环内的train_dataloader改为参数传入的dataloader，同时修正loss计算的平均逻辑：

def train_one_epoch(
    model: nn.Module,
    dataloader:  torch.utils.data.DataLoader,
    loss_fn: nn.MSELoss,
    optimizer: optim.Adam,
    epoch_index,
    scaler: amp.GradScaler, 
):
    model.train()
    running_loss = 0.0
    last_loss = 0.0
    batch_index = 0

    for batch, loader in enumerate(dataloader):
        input_img = loader['input'].to(device, non_blocking=True)
        gt_img = loader['gt'].to(device, non_blocking=True)

        model.zero_grad(set_to_none=True)
      
        with amp.autocast():
            output = model(input_img)
            loss = loss_fn(output, gt_img) 
       
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

        running_loss += loss.item()
        if batch_index % 100 == 99:
            last_loss = running_loss / 100  # 100个batch取平均，修正之前的除以1000错误
            running_loss = 0.0
        
        batch_index += 1

    return last_loss

3. 额外优化建议 #

• 显式初始化权重：为卷积层指定适配LeakyReLU的初始化，强化梯度流动：

  def init_weights(m):
      if isinstance(m, nn.Conv2d):
          nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='leaky_relu')
          if m.bias is not None:
              nn.init.constant_(m.bias, 0)
  model.apply(init_weights)
  

• 学习率调度：建议从1e-4开始设置学习率，配合ReduceLROnPlateau调度器，在loss停滞时自动降低学习率。
• 数据一致性检查：确认GT图像是否同步归一化到[0,1]区间，若GT为[0,255]，需统一归一化后再计算loss。

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