稀疏振幅回归任务中CNN输出恒定值，无法学习真实像素值

问题背景 #

我正在训练一个小型CNN，用于从二值掩码预测稀疏振幅图：

• 输入：60×60的图像，仅15个像素为1，其余为0
• 目标：同尺寸图像，对应输入中1的像素值在0.6-1.0之间，其余为0
• 数据规模：1800张训练图，400张测试图

模型出现输出坍缩问题：所有预测值均为恒定值，尽管目标值存在差异。已尝试的调优手段包括：

• 将振幅缩放至-5到5、-3到3、-1到1区间，测试时反缩放
• 使用Adam、AdamW不同优化器
• 采用SmoothL1Loss、MSELoss不同损失函数
• 调整epoch数和学习率
• 对每个像素单独计算MSE而非整体计算

对比发现：相同架构用于相位预测（目标值范围-π到π，输入与目标位置对应）时学习效果极佳，证明CNN具备学习能力，但振幅任务存在核心问题——输入为1.0，目标为0.6-1.0的真实振幅值，模型无法从输入中获取区分样本的特征。

核心原因分析 #

1. 输入特征缺失区分度：相位任务中，输入掩码的位置直接对应目标相位值，模型能学习位置到相位的映射；但振幅任务中，所有有效输入像素都是1，没有额外特征能关联到不同样本的振幅差异，模型无法从单一值中学习到波动的目标。
2. 损失函数被零值主导：输入和目标中99%以上的像素都是0，默认MSE对所有像素平等计算，有效像素的损失占比极低，模型只需输出恒定值（比如目标均值）就能获得较低的整体损失，自然陷入坍缩。
3. 模型输出范围无约束：振幅目标在0.6-1.0区间，但当前模型最后一层无激活函数，输出范围不受限，容易向极端值或均值坍缩。

针对性解决方案 #

1. 重构输入，加入样本特异性特征 #

既然掩码本身无法区分振幅，需要给模型提供能关联到振幅的额外信息：

• 如果有样本元数据（比如生成掩码的参数），将其作为额外通道加入输入（比如新增一个全图为该元数据值的通道，变成2通道输入）。
• 加入位置编码：对每个值为1的像素，将其归一化坐标（x/60, y/60）作为额外通道，让模型学习“不同位置组合对应不同振幅”的模式。

2. 加权损失，聚焦有效像素 #

修改损失计算逻辑，仅对目标中非零的像素计算损失，避免零值稀释梯度信号：

def weighted_mse_loss(pred, target):
    # 生成有效像素掩码
    mask = (target != 0).float()
    # 仅计算有效像素的MSE，避免被零值主导
    loss = torch.sum(mask * (pred - target)**2) / torch.sum(mask)
    return loss

在训练函数中替换原有的criterion = nn.MSELoss()即可。

3. 约束模型输出范围 #

在模型最后一层加入激活函数，将输出限制在目标区间内：

class SingleOutputFCN(nn.Module):
    def __init__(self):
        super(SingleOutputFCN, self).__init__()
        self.net = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, padding='same'),
            nn.LeakyReLU(0.1),
            nn.Conv2d(32, 32, kernel_size=3, padding='same'),
            nn.LeakyReLU(0.1),
            nn.Conv2d(32, 1, kernel_size=3, padding='same'),
            nn.Sigmoid()  # 将输出压缩到0-1
        )

    def forward(self, x):
        pred = self.net(x)
        # 缩放至目标区间0.6-1.0
        return pred * 0.4 + 0.6

同时可以初始化最后一层的偏置为目标均值（比如0.8），让模型初始输出接近合理范围。

4. 提升模型容量与训练策略 #

• 增加模型层数或通道数：当前仅2层卷积，容量不足，可新增1-2层Conv2d+LeakyReLU，或把通道数从16改为32，提升学习复杂映射的能力。
• 使用学习率调度器：比如torch.optim.lr_scheduler.ReduceLROnPlateau，当损失不再下降时自动降低学习率，避免模型陷入局部最优。

5. 数据增强与样本加权 #

• 对有效像素位置做微小扰动（±1像素），增加数据多样性，帮助模型学习更鲁棒的映射。
• 如果不同振幅区间的样本数量不平衡，给稀有区间的样本赋予更高权重，避免模型偏向预测占比高的振幅值。

完整训练代码修改示例（关键部分） #

def train(model, dataloader, device, epochs=10, lr=1e-3):
    model.to(device)
    model.train()
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=2, factor=0.5)
    
    def weighted_mse(pred, target):
        mask = (target != 0).float()
        loss = torch.sum(mask * (pred - target)**2) / torch.sum(mask)
        return loss

    for epoch in range(epochs):
        total_loss = 0
        for inputs, targets in dataloader:
            inputs = inputs.to(device)
            targets = targets.to(device)

            optimizer.zero_grad()
            pred = model(inputs)
            loss = weighted_mse(pred, targets)
            loss.backward()
            optimizer.step()

            total_loss += loss.item()

        avg_loss = total_loss / len(dataloader)
        print(f"[AMPL] Epoch {epoch+1}/{epochs} | Avg Loss: {avg_loss:.6f}")
        scheduler.step(avg_loss)

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