PyTorch训练CIFAR10自编码器输出黑白图像问题求解

现有实现梳理 #

网络结构 #

class OurAE(nn.Module):
    def __init__(self, in_channels, z_channels):
        super(OurAE, self).__init__()
  
        self.tot_diff = None
        self.in_channels = in_channels
        curr_channels = in_channels
    
        encoder = []
        channels = [3, 16, 64] + [z_channels]
    
        for out_channels in channels:
            encoder += [
                nn.Conv2d(in_channels=curr_channels, out_channels=out_channels, kernel_size=3, padding=1, stride=2),
                nn.ReLU()
            ]
            curr_channels = out_channels
    
        self.encoder = nn.Sequential(*encoder)
    
        curr_channels = z_channels
        decoder = []
        channels = [64, 16, 3] + [in_channels]
    
        for out_channels in channels:
            decoder += [
                nn.ConvTranspose2d(in_channels=curr_channels, out_channels=out_channels, kernel_size=4, padding=1, stride=2),
                nn.ReLU()
            ]
            curr_channels = out_channels
    
        decoder = decoder[:-1] # 移除最后一层ReLU激活
        decoder.append(nn.Sigmoid())
    
        self.decoder = nn.Sequential(*decoder)
        
    def forward(self, x):
        return self.decoder(self.encoder(x))

训练配置 #

• 损失函数：torch.nn.CrossEntropyLoss()
• 优化器：torch.optim.Adam(classifier.parameters(), lr=lr)
• 已尝试调试：将网络最后一层的Sigmoid激活替换为ReLU，问题无改善
• 观测现象：重建得到的输出图像全部为无彩色的黑白灰度样式

核心问题定位 #

• 训练目标完全错配
  交叉熵损失是专为分类任务设计的，用于匹配模型输出的类概率和真实标签，和自编码器需要完成的像素级重建任务毫无适配性。自编码器的训练目标是让输出像素值与输入像素值尽可能接近，正确的损失应该选择像素级回归损失，比如nn.MSELoss()或者nn.L1Loss()。用分类损失训重建任务，模型根本无法对齐重建目标，输出自然是无意义的灰度结果。
• 优化器绑定参数完全错误
  优化器当前绑定的是classifier.parameters()，但当前训练流程中根本没有实例化分类器模块，优化器完全无法获取到自编码器的可训练参数，意味着整个训练过程中自编码器的权重根本没有被更新，相当于随机初始化的网络直接输出结果，三个通道的输出值高度趋同，自然在视觉上呈现为无颜色的黑白灰度图，根本不可能得到正常的彩色重建图像。
• 结构逻辑存在隐性bug
  编码器初始化时，curr_channels直接取值为in_channels，而通道列表的第一个值就是3（即CIFAR10的输入通道数），导致第一次卷积的输入和输出通道完全一致，首次下采样阶段没有做任何通道维度的特征变换；同时潜变量层之后接的ReLU激活会把所有负数值压为0，大幅压缩潜变量的表达范围，即便修正了上面两个硬错误，也容易导致重建结果丢失彩色细节。

快速验证修复步骤 #

• 第一优先级修正：将损失函数替换为nn.MSELoss()，训练时的输入和目标都使用原始的自编码器输入图像，不要接入分类任务的one-hot标签。
• 第二优先级修正：把优化器的参数来源从classifier.parameters()改成实例化的自编码器模块的参数，比如torch.optim.Adam(ae_model.parameters(), lr=1e-3)。
• 可选优化：去掉潜变量层后的ReLU激活，给潜变量保留正负值的表达空间，提升细节重建能力。

内容的提问来源于stack exchange，提问作者Dr. Prof. Patrick