首页 / 资讯中心 / 文章详情

AAE结合ResNet失败复盘:对抗训练不稳定的原因与排查思路

AAE结合ResNet失败复盘:对抗训练不稳定的原因与排查思路 ★ FEATURED ARTICLE
先说结论这个项目最终没有在我的计划时间内跑通。任务是搭建一个AAEAdversarial Autoencoder对抗自编码器我决定不按最常见的套路走而是把生成器从普通卷积堆叠换成了ResNet式的残差块结构。结果挺标准的重建图像模糊到没法看隐空间分布直接坍缩训练loss忽上忽下。前后折腾了差不多一周最后换了方案才稳住。这篇文章就是这次失败项目的完整复盘。说老实话把失败过程写出来需要一点心理建设但这类踩坑记录在生成模型这个领域往往比成功案例更有参考价值。因为AAE本身是个相对小众的方向能查到的资料大多是原理介绍和标准实现真正讲“ResNet结构塞进AAE时为什么崩、崩在哪、怎么查”的内容很少。所以我这篇既是一份技术笔记也算一份排错手册。适合正在折腾AAE、想给生成器换骨干网络、或者遇到对抗训练不稳定问题的朋友参考。1. 项目背景与核心动机1.1 AAE这个结构到底在做什么先花点篇幅把AAE的机制讲清楚因为后面所有失败现象都跟它的结构特点强相关。AAE全称Adversarial Autoencoder本质是一个自编码器额外加了一个“隐变量判别器”。普通自编码器只有一个目标输入图片x经过编码器E压缩成隐向量z再经过解码器D重建出x̂希望重建误差尽可能小。损失函数就一个L_rec ||x - x̂||_1 或者 ||x - x̂||_2^2但这样训练出来的隐空间是没有结构的——z被编码成什么样全看网络自己发挥你不知道它的分布长什么样也就没法从里面采样来生成新图。VAE的解法是强制z服从标准正态分布用的是KL散度约束AAE换了个思路在隐空间加一个判别器让编码器输出的z和从标准正态分布采样的z分别作为假样本和真样本训练判别器区分它们反过来促使编码器把z推向正态分布。所以AAE至少有三个损失在打架重建损失保证解码结果像原图。隐变量对抗损失保证z的分布逼近先验。如果还加了图像域判别器那就是第三个对抗损失用来提升重建/生成图像的清晰度。我当时用的是三者全开的版本。现在回头看这就是埋雷的开始——三个目标之间天然存在竞争关系稍有不平衡就会崩。1.2 为什么要把ResNet塞进生成器既然AAE的编码器负责特征提取解码器负责重建那么很自然的一个想法是编码器用ResNet解码器也堆ResBlock两边都吃残差结构的红利。我当时有三点理由残差连接理论上能让梯度更顺。分类任务里ResNet把梯度消失问题解决得很漂亮生成任务同样要跨很多层反向传播按理说也能受益。恒等映射可以保留低频信息。生成图像最怕把整体色调、轮廓丢掉残差块里那条shortcut路径相当于给了网络一个“原样传递”的选项对重建应该友好。复用现成组件。工程上最现实的原因就是手头已经有ResBlock的实现想直接拿去生成器用省事。听起来很美好。但我忽略了一个关键问题ResNet在分类任务里的成功是建立在“大量监督信号 判别式目标任务”之上的而生成模型里的“监督”只有一个重建loss对抗loss只会给出一个模糊的“像不像”信号。把ResNet搬到生成器里它要面对完全是另一套游戏规则。这一点在后来的排查里被反复验证。2. AAE结合ResNet的整体架构设计2.1 编码器先下采样再堆残差块编码器部分我保留了常见的AAE结构输入是3通道32x32图像先经过几个带stride的卷积快速下采样把空间分辨率压到4x4然后再接几个残差块做非线性变换最后拉平过全连接层输出128维隐变量。核心代码大概是这样的PyTorch风格import torch import torch.nn as nn import torch.nn.functional as F class ResBlock(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.conv1 nn.Conv2d(in_ch, out_ch, 3, 1, 1) self.bn1 nn.BatchNorm2d(out_ch) self.conv2 nn.Conv2d(out_ch, out_ch, 3, 1, 1) self.bn2 nn.BatchNorm2d(out_ch) self.shortcut nn.Identity() if in_ch out_ch else nn.Conv2d(in_ch, out_ch, 1) def forward(self, x): identity self.shortcut(x) out F.relu(self.bn1(self.conv1(x))) out self.bn2(self.conv2(out)) return F.relu(out identity) class Encoder(nn.Module): def __init__(self, latent_dim128): super().__init__() self.stem nn.Sequential( nn.Conv2d(3, 64, 3, 2, 1), nn.BatchNorm2d(64), nn.ReLU(True), nn.Conv2d(64, 128, 3, 2, 1), nn.BatchNorm2d(128), nn.ReLU(True), nn.Conv2d(128, 256, 3, 2, 1), nn.BatchNorm2d(256), nn.ReLU(True), ) self.resblocks nn.ModuleList([ResBlock(256, 256) for _ in range(4)]) self.fc nn.Linear(256 * 4 * 4, latent_dim) def forward(self, x): h self.stem(x) for block in self.resblocks: h block(h) h h.view(h.size(0), -1) return self.fc(h)这里值得说明的是为什么要添加stem下采样直接在32x32的输入上跑残差块计算量会很大而且浅层的高频细节对隐变量压缩没有帮助先降下来再提取高层语义是自编码器常用的做法。2.2 生成器低分辨率空间狂堆残差块生成器也就是AAE里的解码器是这个项目里最“上头”的设计。我的思路是隐变量z先过全连接层映射成4x4特征图然后接一堆ResBlock在低分辨率空间里做非线性变换最后用几个转置卷积把分辨率拉回32x32。class Generator(nn.Module): def __init__(self, latent_dim128): super().__init__() self.fc nn.Linear(latent_dim, 256 * 4 * 4) self.resblocks nn.ModuleList([ResBlock(256, 256) for _ in range(8)]) self.upsample nn.Sequential( nn.ConvTranspose2d(256, 128, 4, 2, 1), nn.BatchNorm2d(128), nn.ReLU(True), nn.ConvTranspose2d(128, 64, 4, 2, 1), nn.BatchNorm2d(64), nn.ReLU(True), nn.ConvTranspose2d(64, 3, 4, 2, 1), nn.Tanh(), ) def forward(self, z): h self.fc(z).view(-1, 256, 4, 4) for block in self.resblocks: h block(h) return self.upsample(h)8个ResBlock在4x4特征图上做处理空间分辨率始终不变通道数也不变等于是在一个很小的特征图上做“深度非线性变换”。我当时觉得这很酷因为主流做法是在上采样过程中不断加卷积层而我这种是在最底层拼命加深度。结果证明这个决策踩中了两个大坑一是深层网络在对抗反馈下极其不稳定二是低分辨率特征图承载不了太深的非线性变换——信息在逐层卷积中不断被扰动重建细节越到后面越丢。2.3 两个判别器与损失函数设计AAE在标准版本里只需要一个隐变量判别器但为了让重建出来的图像更清晰我额外加了一个图像域判别器。隐变量判别器是标准的MLP判断输入的128维向量是来自先验分布还是来自编码器输出。图像判别器用的是PatchGAN结构输出一个N x N的矩阵对图像的每个局部区域判断真假。class LatentDiscriminator(nn.Module): def __init__(self, latent_dim128): super().__init__() self.net nn.Sequential( nn.Linear(latent_dim, 256), nn.LeakyReLU(0.2), nn.Linear(256, 256), nn.LeakyReLU(0.2), nn.Linear(256, 1), ) def forward(self, z): return self.net(z)损失函数采用LSGAN的形式避开标准GAN的饱和问题重建损失L1距离权重10隐变量对抗损失LSGAN权重1图像对抗损失LSGAN权重1总损失是三者加权求和。这个权重我后来调过很多次但不管怎么调训练都撑不过30轮。核心问题不在权重微调而在结构本身。3. 训练过程实录与失败现象3.1 训练配置与基线预期数据集用的是公开图像数据集32x32分辨率规模不算大。训练配置如下配置项实验值batch size32优化器Adam学习率2e-4beta(0.5, 0.999)隐变量维度128生成器ResBlock数量8编码器ResBlock数量4训练轮数100重建loss权重10图像对抗loss权重1隐变量对抗loss权重1打开训练日志前我的心理预期是重建loss会先降然后隐变量分布逐步逼近正态分布图像质量在中后期稳定下来。另外因为AAE同时有重建和对抗两个目标收敛标志应该是重建图像清晰、采样生成图像也像样。3.2 训到二十轮左右的异常曲线前5个epoch重建loss从1.5左右快速降到0.4附近生成出来的图像虽然模糊但轮廓是对的。这个阶段给了我很大信心以为一切尽在掌握。大概从第8个epoch开始情况不对了。重建loss降到0.35左右之后就不再下降生成图像进入一个“模糊平台期”怎么都清晰不起来。最让我困惑的是两个对抗loss的表现非常反差图像判别器的loss掉到了0.3以下但生成器的对抗loss却在20到80之间剧烈震荡。这种“判别器赢麻了、生成器输惨了”的状态在GAN里见的多了通常意味着模式坍缩或者判别器过强。到了第20个epoch左右重建loss开始不降反升生成图像的细节彻底崩坏。有些样本变成纯色块有些出现棋盘格伪影还有的输出图像里同一通道全是同一个值。训练日志里的loss曲线就像心电图已经不是“波动”而是“蹦极”。3.3 三个肉眼可见的失败信号梳理一下当时观察到的具体现象这些信号在生成模型训练里非常典型隐变量坍缩。把编码器输出的z拉出来做了个统计发现128个维度的方差从接近1掉到了0.01以下。意味着编码器输出的z几乎退化成一个常数不管输入什么图中间向量都差不多。这是最致命的信号——隐变量判别器的对抗压力没把z推向正态分布反而把编码器逼到了“躺平”状态。重建图像出现“死亡通道”。生成图像里某些通道完全恒定说明生成器的部分卷积核在训练中失效输出被激活函数压到了恒值。这个问题在深层生成器里尤其明显梯度在几十层间反复传播后部分路径的权重更新已经停止。训练loss不可复现。同一个配置重启一次训练前面几个epoch的loss趋势相同但崩掉的时间点不同——有时候12轮崩有时候18轮崩。这正是对抗训练不稳定性的典型特征说明问题不是单纯的数据或超参问题而是整个训练博弈在结构上缺乏约束。3.4 失败模式速查表为了后面排查方便我把当时的失败现象整理成了下面这张表失败现象可能原因判断依据重建loss长期停在0.35不下降解码器容量不足或梯度路径失效纯自编码器模式可以继续下降隐变量方差趋近于0编码器在对抗压力下坍缩z统计值方差远小于1生成器对抗loss剧烈震荡判别器过强或生成器梯度不稳定loss在20-80之间跳变通道恒定、色块伪影深层卷积核失效与转置卷积叠出棋盘格特征图可视化发现恒定通道同配置不同轮次崩溃对抗训练缺乏稳定约束多次重启实验验证现在回看这张表其实已经能看出问题的大致方向了但当时处在“跑实验、看loss、调参、再跑实验”的循环里没有跳出来做系统性拆解。4. 问题排查与原因定位4.1 排查流程先隔离再组合如果一上来就同时怀疑结构、损失、超参那肯定调到头秃也找不到根因。我的做法是一步一步做减法每次只改变一个变量看能不能复现或者消除某个现象。排查步骤分四步走第一步把两个判别器全部关掉只留编码器-解码器和重建loss。此时模型变成一个纯自编码器。结果重建loss可以降到0.1以下图像轮廓和颜色都恢复得不错。这个结论很重要编码器和生成器本身的容量是够的问题出在加入对抗训练之后。第二步只开隐变量判别器关掉图像判别器。重建质量开始变差z的分布统计也出现坍缩趋势。说明隐变量对抗训练对编码器的负面影响是独立存在的。第三步只开图像判别器关掉隐变量判别器。图像锐度有所提升但训练稳定性变差后期还是会崩。说明图像对抗也引入了不稳定因素。第四步把生成器的ResBlock数量从8降到3两个判别器全开。训练能撑过50轮但生成图像细节差重建结果依然不如纯自编码器。通过这一轮隔离实验基本可以断定ResBlock堆叠的深度不是唯一问题隐变量对抗和图像对抗各有各的雷叠加起来更是灾难。4.2 根因一深层ResBlock在生成任务中的信息损耗ResNet的分类成功让很多人误以为“更深更好”但这个前提是模型在高维特征空间中做判别。生成任务的本质是“从压缩表示还原细节”它需要的是信息保真而不是特征抽象。我在4x4特征图上堆8个ResBlock每个Block的卷积-归一化-激活都会对特征做一次非线性变换。恒等映射那条shortcut路径虽然能传递原始信息但主路径的扰动经过层层叠加会在重建阶段被放大。尤其是对抗loss加入后梯度信号变得很噪残差块的恒等路径很难维持“原样传递”信息被逐层抹掉。这个根因后来用一个简单实验验证了把生成器里的ResBlock全部去掉改成三个普通卷积直接上采样重建loss反而降得更低。深度在这里不仅没有帮忙还在帮倒忙。4.3 根因二BatchNorm在对抗生成网络里的放大效应BatchNorm在分类任务里是好助手它用batch内的均值和方差对特征做归一化。但在生成模型里它经常成为不稳定源。原因是这样的判别器和生成器在对抗训练中会不断改变数据分布batch的统计量因此发生剧烈波动。比如一个batch里恰好生成样本特别“假”那么这一轮的BN均值和方差就会偏向某一侧迭代之间的统计量抖动被放大。这种抖动传递到梯度里会让生成器参数更新步长忽大忽小表现出来的就是loss剧烈震荡。我在排查时做了一次对照实验把生成器里的BatchNorm全部换成InstanceNorm重建loss的稳定性明显改善。InstanceNorm对单张样本做归一化不受batch内其他样本影响在生成任务里通常比BN稳得多。这个结论在GAN社区其实早有共识但我当时没放在心上硬是踩了一遍才长记性。4.4 根因三对抗目标诱发隐变量坍缩AAE里隐变量判别器的设计本意是“迫使编码器输出的z服从标准正态分布”。但对抗训练是零和博弈编码器作为一个参与者它有强烈的动机寻找“欺骗判别器的最短路径”。对编码器来说最短路径不是把z分布逼近正态分布而是把所有x都映射到z空间的一个小区域——最好是一个点附近。这样判别器无论如何都分不清真假因为它看到的z都来自同一小片区域方差足够小看起来“像”一个聚集的先验样本但实际上是坍缩的。想通过单纯加大隐变量对抗权重来解决是无效的因为更大的对抗压力只会让编码器更努力地寻找这个捷径。真正有用的做法是给隐变量加正则或者在重建loss上加大筹码编码器如果敢把z压缩成一个常数重建loss就会飙升它会因为重建目标而被迫保持z的信息量。5. 可行的修复思路与替代方案5.1 第一刀给残差块瘦身如果一定要坚持用ResNet做生成器不要像我一样堆8个Block。把数量降到2到3个并且用BasicBlock而不是Bottleneck——生成任务的特征图通道数通常不会超过512Bottleneck的降维-升维结构反而会丢细节。我当时把ResBlock数量降到3个之后训练稳定性有明显改善。如果再配合残差块内部的“预激活”结构也就是把BN和ReLU放到卷积之前还可以进一步减小主路径扰动的累积。5.2 第二刀替换归一化层在生成器里优先使用InstanceNorm或GroupNorm不要用BatchNorm。判别器倒是可以用BatchNorm但更好的做法是给判别器加谱归一化约束它的Lipschitz常数避免判别器梯度变得过大。我当时把生成器的BN换成InstanceNorm之后生成器loss的震荡幅度小了一个量级。替换归一化层是这次失败排查中性价比最高的一项调整。5.3 第三刀换掉转置卷积的上采样方式转置卷积叠多了会产生棋盘格伪影这是生成模型里出了名的坑。我的生成器用了三个转置卷积每个stride2输出分辨率从4x4一路拉到32x32棋盘格问题非常明显。替代方案是先用最近邻插值或者双线性插值放大分辨率再接普通卷积做细化或者用PixelShuffle做亚像素卷积。我当时试了PixelShuffle棋盘格伪影立刻减轻重建图像的纹理也更自然。5.4 第四刀重构损失与训练策略损失函数的配比需要重新考虑。我最初的设置是L1重建权重10、图像对抗1、隐对抗1这个比例在浅层模型上勉强能跑但在深层模型上根本压不住。更稳的策略分两个阶段阶段一只训练自编码器用重建loss跑20到30个epoch让编码器和生成器先达到一个不错的重建水平。这一步相当于预训练。阶段二冻结一个合理的初始状态后再同时开启两个判别器。判别器的学习率要比生成器低一些比如2e-5生成器和编码器保持2e-4。这样对抗训练是在一个稳定的重建基础上“微调”而不是从零开始对抗博弈。在这个策略下隐变量坍缩问题会缓解很多因为编码器在阶段一里已经习惯了输出有信息的z阶段二的对抗压力很难把已经学会的表示完全摧毁。最终我参考的稳定版配置是配置项调整后数值生成器ResBlock数量3生成器归一化层InstanceNorm上采样方式PixelShuffle Conv图像判别器PatchGAN 谱归一化预训练轮数20对抗训练阶段lr生成器2e-4判别器2e-5重建loss权重10图像对抗loss权重0.5隐变量对抗loss权重0.5这套配置跑下来至少能把训练过程镇住重建质量明显好于初始版本。当然它不是最优解只是给同样想折腾AAE结构的朋友一个不崩的起点。6. 复盘与经验教训6.1 深度不是生成质量的保证这次失败让我对“深度”有了新的理解。ResNet在判别式任务里越深越强是因为有清晰的目标函数和足够多的监督信号但在生成模型里监督信号极度稀疏重建loss只能告诉你“像不像”对抗loss只能告诉你“真不真”。没有更细粒度的信号约束深层网络很容易在训练中摆烂——这不是网络的错是任务本身没有给网络足够的方向。6.2 拆解验证比盲目调参更重要回头看整个排查过程最浪费时间的就是在前面反复调超参的阶段。什么把学习率从2e-4降到1e-5、把batch size从32改成64、给生成器加dropout……全都试了一遍没有一个能解决根本问题。直到我开始做隔离实验把两套对抗训练全部关掉才意识到结构本身有病。以后再做类似的模型组合我一定会先跑通最简版本再逐步加复杂度。每一步都确认“加了这个东西之后loss有没有崩效果有没有提升”全部验证通过才继续往下走。6.3 失败模式库是黑暗里的一盏灯训练中遇到的那些“death channel”、“隐变量坍缩”、“判别器太强”其实都是生成模型社区很常见的问题并不是什么新鲜事。只是当时手头没有一份系统性的失败模式对照表只能靠自己猜。现在我把这次的经历整理成文也是在给自己攒一个排错手册。如果大家在做类似方向时遇到同样的现象至少能少走几轮弯路。6.4 给后来者的三句话如果让我把这个项目重新做一遍我会按照下面的优先级来安排资源先把纯自编码器在目标数据集上训练到重建满意再加任何对抗模块。生成器的骨干优先选稳定结构——浅层残差、InstanceNorm、PixelShuffle而不是盲目堆Block。一旦训练开始崩先做隔离实验定位是哪个损失在破坏整体不要一上来就怀疑超参。这个东西后续如果想继续扩展可以考虑引入感知损失来约束重建的语义内容或者用变分自编码器和AAE做混合给隐变量一个显式的分布监督。不过这些都是后话了先把基础版本稳住才是这个方向最值得投入的事情。
阅读完成 · 觉得有帮助?
咨询建站