预训练ResNet-50作为CAE编码器性能不及简易CAE的原因咨询

问题背景 #

我目前正在测试不同卷积自编码器（Convolutional Autoencoder, CAE）架构，尝试将在ImageNet上预训练的ResNet-50作为模型编码器，共测试了两种方案：一是固定编码器权重不更新，二是将预训练权重作为初始化参数参与训练。后者重建效果更优，但仍不及编码器含3层卷积、解码器含3层上采样+Conv2d层的基础CAE。我在OfficeHome数据集的不同域上训练模型，相关架构代码如下：

class ConvBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(ConvBlock, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1)
        self.relu = nn.LeakyReLU(inplace=True)
        self.norm = nn.BatchNorm2d(out_channels)
        
    def forward(self, x):
        x = self.conv(x)
        x = self.relu(x)
        x = self.norm(x)
        return x

class ResnetConvAutoencoder(nn.Module):
    def __init__(self):
        super(ResnetConvAutoencoder, self).__init__()
        resnet = models.resnet50(pretrained = True, progress = True)
        
#         for param in resnet.parameters():
#             param.requires_grad = False
        
        self.encoder = torch.nn.Sequential(*(list(resnet.children())[:-2])) #output is 2048 x 8 x 8

        del resnet

        self.decoder = nn.Sequential(
           nn.Upsample(scale_factor=2, mode='bilinear'),
           ConvBlock(2048, 512),
           
           nn.Upsample(scale_factor=2, mode='bilinear'),
           ConvBlock(512, 256),
            
           nn.Upsample(scale_factor=2, mode='bilinear'),
           ConvBlock(256, 128),
            
           nn.Upsample(scale_factor=2, mode='bilinear'),
           ConvBlock(128, 16),
           
           nn.Upsample(scale_factor=2, mode='bilinear'),
           ConvBlock(16, 3),
          
           nn.Sigmoid()
        ) # output is 3 x 256 x 256 (as input)
        
        self.encoder.cuda(2) # trainig on several GPUs
        self.decoder.cuda(1)

    def forward(self, x):
        x = x.cuda(2)
        x = self.encoder(x).cuda(1)
        x = self.decoder(x).cuda(1)
        return x

重建效果对比 #

• ResNet-50编码器CAE重建结果：
• 低层数简易CAE重建结果：

训练及测试阶段，ResNet-50作为编码器的CAE的MSE损失稳定在0.03左右无法下降，而自研简易CAE的MSE损失可低于0.015。

性能差异原因分析 #

• 预训练特征适配性差：ImageNet预训练的ResNet-50是为分类任务设计，学习目标是区分不同类别主体，因此会刻意压缩、丢弃像素级的纹理、边缘等底层细节信息，仅保留高层语义特征。而自编码器重建任务高度依赖多尺度的细节信息，ResNet-50编码器输出的8x8高维特征本身就缺失大量重建所需的底层信息，天然限制了重建效果的上限。
• 解码器架构不合理：当前使用的ConvBlock全部采用1x1卷积，仅能完成通道维度的数值变换，无法融合空间邻域的特征信息。仅依靠双线性上采样+1x1卷积的解码器结构，没有能力填补32倍下采样带来的信息损失，而常规的3层简易CAE通常采用3x3卷积，具备空间特征融合能力，更适配图像重建任务。
• 下采样倍率差距大：输入为256x256分辨率时，ResNet-50编码器的下采样倍率为32倍，输出特征图尺寸仅为8x8；而3层卷积的基础CAE下采样倍率通常仅为8倍，输出特征图尺寸更大，保留的空间细节信息更多，重建难度更低，自然可以取得更低的MSE损失。
• 模型训练不充分：ResNet-50编码器的参数量远大于3层卷积的简易编码器，整体模型参数量级存在明显差距，需要更多的训练数据、更长的训练轮次才能完全收敛。OfficeHome数据集规模有限，很容易出现模型训练不充分的问题，导致损失卡在0.03无法继续下降。
• 跨卡部署的间接影响：将编码器和解码器部署在不同GPU上，每次前向传播都需要跨卡拷贝数据，若训练时batch size设置较小，可能会间接影响梯度更新的稳定性，拉长收敛所需的轮次。

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