读 StarGAN 的源码之前,我建议先把论文里那句"一个生成器搞定多个域"放在一边,因为真正决定这套代码能不能跑通、能不能改出自己数据的,是张量在几个关键节点上的形状变化、标签在什么时候被拼进输入、以及判别器那两个输出头到底谁负责什么。我前后把这份代码翻过好几遍,第一遍是照着跑通,第二遍是为了把它迁到自己的一套人脸属性数据上,第三遍干脆是为了给团队里新来的同学做讲解。三次下来感受很一致:StarGAN 代码解析的难点从来不在网络有多深,而在多域翻译这件事被拆成了一堆标签拼接和损失配比的细节,任何一个地方对不上,损失曲线就会给你一个看起来在降、实际根本没学到东西的假象。
这篇内容适合两类人看。一类是已经把条件生成模型跑通过、想弄清楚"多域"到底怎么落到代码层的人;另一类是拿到了类似的多属性编辑需求,想知道单生成器多域这条路值不值得走、坑在哪儿的人。下面我按自己读代码时的真实顺序来讲:先弄清它要解决什么结构性问题,再一层层拆生成器、判别器、损失函数和训练循环,最后说说我怎么验证它是不是真的学到了东西。所有代码片段都是我按原实现的结构重写的示意版本,变量命名做了统一,方便你对照自己的仓库。
1. 多域翻译的结构性难题:为什么不能简单地堆多个生成器
1.1 一对多模型堆叠带来的参数冗余与标签割裂
最直觉的做法是每个域对训一个生成器。假设你要做人脸的五种属性编辑,头发颜色、性别、年龄、表情、肤色,两两组合就有二十个方向,去掉自映射还剩二十个模型。每个模型一套参数、一套训练流程、一套超参,显存和存储先不说,光是数据组织就够喝一壶——同一张图要按二十种配对方式反复喂进去,每换一个方向就重训一次。
更麻烦的是标签割裂。每对独立的模型只认识自己那两个域,模型 A 知道"黑发到金发",模型 B 知道"男到女",但它们之间没有任何共享表示。你想同时改发色和性别,只能串行调用两个模型,第二次调用会把第一次的结果再当输入处理一遍,伪影叠加得非常明显。代码层面这表现为一个非常尴尬的调度问题:串行链式调用时,中间结果要不要做后处理?不做好像也能看,做好像又引入了额外的不确定性。
StarGAN 的核心主张就是把这些模型压成一个:用同一个生成器,靠输入里的域标签决定往哪个方向翻译。这样参数只有一份,任意两个域之间的翻译共享同一套骨干特征,多属性同时编辑也变成了一次前向,标签向量里同时置多个位就行。这个改动的价值不在于省了多少显存,而在于它把"域"从一个模型级别的概念降到了数据级别的概念。
1.2 域标签作为输入的一部分:这个设计决定的连锁反应
把标签拼进输入看起来是个小技巧,但它连锁决定了好几件事。生成器的第一个卷积层输入通道数不再是 3,而是 3 加上标签维度;标签必须是连续可微的数值向量而不是类别索引,否则没法参与反向传播;标签的空间维度必须被广播到和特征图一样的尺寸,所以实际代码里通常是把标签向量在空间维度上复制成 H×W 的常数图再拼接。
这带来一个很实际的编码陷阱。很多人第一次写的时候会把标签拼成形状(N, c_dim)然后直接和(N, 3, H, W)的图像做torch.cat,运行时报维度不匹配。正确做法是先view(N, c_dim, 1, 1)再expand到(N, c_dim, H, W)。这个 expand 不做真实内存拷贝,只改 stride,所以开销可以忽略,但如果你用repeat就会实实在在占一份显存。
1.3 掩码向量存在的意义:多数据集联合训练的必要件
原实现里还有一个 mask vector 的东西,很多人在单数据集场景下会觉得多余。它的设计目的很明确:当你同时用好几个数据集训练时,不同数据集支持的属性集合不一样,A 数据集有发色标签但可能没有年龄标签,B 数据集反过来。如果直接把不存在的属性当 0 处理,模型会误以为"这个属性确实是 0"而不是"这个属性未知"。
mask 向量的作用就是显式标注哪些位是有效的。在单数据集场景下,mask 全为 1,这一整套逻辑等于没起作用,所以第一次读代码时完全可以先跳过它,等你要联合多个来源的数据时再回头补上。这一点我后面还会提到,因为它牵扯到损失函数里那个容易被忽略的mask参数。
2. 生成器逐层拆解:从通道拼接到底层特征的重建
2.1 输入构造:拼接顺序和维度对齐的实操细节
生成器的第一层决定了整条链路的形状基准。原实现里,输入是图像和标签在通道维拼接,图像在前、标签在后,通道数为3 + c_dim。这个顺序本身不是强制的,但它必须和推理时保持一致。我曾经为了"代码更整洁"把标签放到前面,训练完发现推理效果全崩,原因就是训练和推理的拼接顺序不一致,模型学到的是"前三通道是图像"这个隐含约定,一换顺序就全乱。
下面这段是生成器的初始化骨架,我按原结构重写并加了注释:
class Generator(nn.Module): def __init__(self, conv_dim=64, c_dim=5, repeat_num=6): super(Generator, self).__init__() layers = [] # 输入通道 = 图像3通道 + 标签c_dim通道 layers.append(nn.Conv2d(3 + c_dim, conv_dim, 7, 1, 3, bias=False)) layers.append(nn.InstanceNorm2d(conv_dim, affine=True)) layers.append(nn.ReLU(inplace=True))c_dim在常见的人脸属性数据上取 5,对应五个二值属性。注意这里是二值属性而不是 one-hot 的互斥类别,这一点和很多多分类的条件生成代码不一样,它允许一个样本同时具有多个属性,这也是后面损失函数用 BCE 而不是交叉熵的原因。
2.2 下采样、残差块与上采样的通道曲线
主体结构是三段式:两个 stride 为 2 的卷积把分辨率降到四分之一,通道翻到四倍;中间接若干个残差块,分辨率不变、通道不变;再用两个转置卷积把分辨率升回去,通道降回来。这套结构在图像翻译任务里被用了很多年,选择它的理由很朴素:低分辨率层负责全局的形变和风格,残差块负责在低分辨率上做足够深的信息交互,上采样负责还原细节。
残差块里我用的是InstanceNorm而不是BatchNorm。这个选择在图像风格类任务里几乎是默认答案,理由是归一化的统计量应该按单张图算,而不是按整个 batch 算。如果你用 BatchNorm,batch 内不同样本的统计量会互相污染,生成结果会出现"跟着同 batch 的其他样本变"的现象。实测下来这个差异在 batch size 小的时候尤其明显,batch size 小于 8 时 BatchNorm 的效果会掉得很厉害。
class ResidualBlock(nn.Module): def __init__(self, dim_in, dim_out): super(ResidualBlock, self).__init__() self.main = nn.Sequential( nn.Conv2d(dim_in, dim_out, 3, 1, 1, bias=False), nn.InstanceNorm2d(dim_out, affine=True), nn.ReLU(inplace=True), nn.Conv2d(dim_out, dim_out, 3, 1, 1, bias=False), nn.InstanceNorm2d(dim_out, affine=True)) self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False) def forward(self, x): return x + self.main(x)self.conv1x1只有在输入输出通道不一致时才会被用到,标准配置下dim_in == dim_out,所以这个分支实际上走的是恒等映射。我在自己改代码的时候把它删掉过,跑起来没区别,但留着更安全,因为你一旦改通道配置就会用到。
2.3 输出层的激活函数与数值边界
最后一层是 7×7 卷积加Tanh。这里没什么可纠结的,输出归一化到[-1, 1]是配套数据预处理决定的。真正容易出问题的是数据侧的归一化方式:如果你的transforms.Normalize用的是 ImageNet 的均值和方差,输出又用 Tanh,那训练目标就是一个不匹配的区间,损失会一直卡在一个不小的值上降不下去。
我踩过这个坑,现象特别有迷惑性:损失在降,但生成的图整体偏色、对比度发灰,看起来像"学了一半"。排查了半天才发现是Normalize用了mean=[0.485, 0.456, 0.406]而输出范围是[-1, 1]。正确做法要么是把Normalize改成mean=0.5, std=0.5,要么输出层换成Sigmoid配套[0,1]区间。这两种都行,但一定要全流程统一,包括推理和可视化的反归一化。
3. 判别器的双头结构:真假与域分类为什么共享骨干
3.1 一条主干两个输出分支的代码组织
判别器是 StarGAN 里最容易被低估的部分。它不是一个简单的真假日判别器,而是在同一套卷积特征上接了两个头:一个头输出单通道,判断真假;另一个头输出c_dim通道,判断这张图属于哪些域。前几层完全共享,直到接近输出才开始分叉。
class Discriminator(nn.Module): def __init__(self, image_size=128, conv_dim=64, c_dim=5, repeat_num=6): super(Discriminator, self).__init__() layers = [nn.Conv2d(3, conv_dim, 4, 2, 1, bias=False), nn.LeakyReLU(0.01, inplace=True)] curr_dim = conv_dim for _ in range(1, repeat_num): layers.append(nn.Conv2d(curr_dim, curr_dim * 2, 4, 2, 1, bias=False)) layers.append(nn.LeakyReLU(0.01, inplace=True)) curr_dim *= 2 self.main = nn.Sequential(*layers) # 真假分支 self.conv1 = nn.Conv2d(curr_dim, 1, 3, 1, 1, bias=False) # 域分类分支,把空间维度压成 1x1 kernel_size = image_size // 2 ** repeat_num self.conv2 = nn.Conv2d(curr_dim, c_dim, kernel_size, 1, 0, bias=False)这里第一个注意力点在repeat_num和image_size的关系上。kernel_size必须满足image_size // 2 ** repeat_num,也就是经过repeat_num次 stride 2 下采样后的特征图尺寸。如果输入是 128×128、repeat_num=6,算出来是 2,那么域分类分支的卷积核就是 2×2,输出 1×1。如果你改了输入分辨率却忘了改repeat_num或者写死 kernel_size,运行时不会报错,但输出形状会不对,损失计算时会崩或者静默出错。
3.2 域分类损失用 BCE 而不是交叉熵的原因
这个分支的输出维度是c_dim,每个位置是独立的一路二值判断。所以损失必须用BCEWithLogitsLoss或者BCE,而不是多分类的交叉熵。原因是这几个属性之间并不互斥——一个人可以同时是"男性"和"戴眼镜"。如果你误用交叉熵,模型的输出会被 softmax 归一化,五个属性的概率和恒为 1,模型被迫在其中选一个最强,另外一个就学不出来。
# 判别器对真实样本的域分类损失 out_src, out_cls = D(real_x) d_loss_cls = F.binary_cross_entropy_with_logits( out_cls, real_label, reduction='none').mean()注意这里real_label是 float 类型的多热向量,形状(N, c_dim),取值 0 到 1。数据加载时要确保转换成了 float,如果是 long 类型会直接报错。这个细节很琐碎但特别常见,尤其是你从分类任务的代码里搬过来的时候。
3.3 两个头共用主干的取舍:省参数还是省不了
有人会问,为什么不干脆训两个网络,一个判真假一个判域。答案在于梯度。如果两个网络分开,判真假的网络只关心纹理和细节是否符合真实分布,判域的网络只关心属性是否可辨,两者学到的特征可能完全不重叠。而共享主干会强制模型找一套"既能判真、又能辨认属性"的特征,这套特征对生成器来说更有约束力。
当然这个设计也有代价。两个损失的梯度会互相干扰,如果两个损失的量级差得多,共享主干就会被其中一个主导。我在调试时见过域分类损失迅速降到接近 0 而对抗损失几乎不动的情况,表现是生成的图属性确实改对了,但图像质量很差、糊。这时候要做的就是调lambda_cls的权重,而不是怀疑网络结构。
4. 损失函数逐项落地:四个目标在代码里的实际配比
4.1 对抗损失:为什么原实现选了带梯度惩罚的形式
原实现采用的对抗损失形式是带梯度惩罚的 Wasserstein 距离(常被简称为 WGAN-GP 那一套)。相较于最原始的 GAN 交叉熵目标,它不需要精心平衡生成器和判别器的学习速度,也不太容易出现梯度消失。代价是每次更新判别器都要额外算一次梯度惩罚项,训练时间会拉长不少。
梯度惩罚的核心是在真实样本和生成样本之间随机插值,然后在插值点上算判别器梯度的范数,约束它接近 1:
alpha = torch.rand(batch_size, 1, 1, 1).to(device) x_hat = (alpha * real_x + (1 - alpha) * fake_x).requires_grad_(True) out = D(x_hat)[0] grad = torch.autograd.grad( outputs=out, inputs=x_hat, grad_outputs=torch.ones_like(out), create_graph=True, retain_graph=True, only_inputs=True)[0] grad = grad.view(batch_size, -1) grad_penalty = ((grad.norm(2, dim=1) - 1) ** 2).mean()create_graph=True这个词很关键。很多人抄代码的时候把它漏了,结果惩罚项虽然算出来了但没有连进计算图,反向传播时这一项的梯度是 0,等于白算。不会有任何报错,只是损失曲线看起来正常但实际没起作用。这种静默失效在 GAN 代码里特别多,我的习惯是每次搭完先打印一下各项损失是否真的在回传(比如检查wgan.grad_penalty.requires_grad)。
4.2 域分类损失:生成器和判别器各算一份,方向相反
域分类损失在两个地方出现。判别器那份是让真实图的属性被正确认出;生成器那份是让生成图的属性被认成目标属性。方向完全相反,但用的是同一个损失函数。
# 生成器:生成图必须被判成"目标标签" g_loss_cls = F.binary_cross_entropy_with_logits( D(fake_x)[1], target_label, reduction='none').mean()这里target_label是生成时喂进去的标签。这个损失是 StarGAN 能做到"属性可控"的核心,没有它,生成器完全可以生成一张高质量但属性不对的图。实操中这一项通常是收敛最快的,往往几百步之后就被判器认得很准了。如果它降不下去,先检查标签是否真的被送进了生成器,以及标签和数据之间有没有错位——数据加载器的 shuffle 和标签数组的索引对不齐,是这类问题的高发区。
4.3 重构损失:把生成结果再变回原图的意义
重构损失的逻辑是:拿原图配上原标签,生成一张翻译图;再拿这张翻译图配上原标签,生成一张"回译图";回译图要和原图尽可能像。这一项约束的是内容保留,防止生成器只学会改属性、把身份信息全丢。
rec_x = G(fake_x, original_label) g_loss_rec = torch.mean(torch.abs(real_x - rec_x)) # L1用 L1 而不是 L2 是个常见选择。L1 对大误差的惩罚不那么极端,生成结果会更锐利一些,不容易糊成一片均值。我在自己数据上对比过,L2 的收敛曲线更平滑,但视觉效果明显更模糊,尤其是头发丝这类细节。
有一个容易忽略的点:这个二次前向是在同一个G上跑的,所以计算图会变大,显存占用比单次前向多不少。如果你显存吃紧,可以降低 batch size,但不要尝试detach中间结果,那样重构损失就没法回传给第一次生成的路径了。
4.4 生成器与判别器的更新顺序与延迟更新策略
原实现里判别器更新n_critic次、生成器更新 1 次。这个比例在 WGAN 类目标里是标配,目的是让判别器先学到足够好的距离估计,再让生成器跟上。我在小数据集上试过把n_critic调到 1,结果是训练前期看起来正常,到中期生成器开始输出大面积的色块,典型的原因是判别器跟不上、给的梯度方向不稳定。
判断n_critic是否合适有个很土但有效的办法:看判别器对真假样本的输出差值。如果两者几乎没差距,说明判别器太弱,需要加大n_critic或者降低生成器学习率;如果差值巨大且持续扩大,说明判别器过强,生成器拿不到有效梯度,要反过来调。
5. 训练主循环里的工程细节:数据、缓存与显存
5.1 数据加载与固定标签采样
训练循环里最容易被忽视的是标签的来源。单域翻译任务里标签是固定的,但 StarGAN 里每次迭代都要重新采一组随机目标标签,这样同一张图在不同的迭代里被推向不同的方向,模型才能覆盖所有域对。
def sample_target(label, c_dim): target = torch.zeros_like(label) idx = torch.randint(c_dim, (label.size(0),)) target[torch.arange(label.size(0)), idx] = 1.0 return target这是个简化版,只采一个随机属性作为目标。如果你要支持多属性同时编辑,可以采样一个伯努利向量,但要注意保证至少有一位为 1,否则目标标签全零,模型学到的就是"把所有属性都关掉",这个空域会污染整个训练。
数据加载层面还有一个隐性开销:属性标签一般跟图片路径一起放在一个列表里一次性读进内存。如果数据集大,这个列表会占不少内存,且加载时要做一次全量遍历。我的做法是提前把标签整理成一个 numpy 文件,训练时用索引读取,启动速度快很多。
5.2 生成样本缓存:为什么要留着上一轮的假图
判别器更新时用到的假样本,在很多实现里会缓存起来,下一轮生成器更新时复用。这么做是为了减少一次前向的计算量。但缓存要小心一点:如果缓存的假图带着上一步生成器的计算图,反传时会出错。必须对缓存的假图做detach()。
fake_x = G(real_x, target_label) fake_x_cached = fake_x.detach() # 给判别器用 # 生成器损失里继续用 fake_x,保留梯度我第一次看这段的时候没理解为什么同一个变量要弄出两份,直到自己写的时候忘了detach,报了一个"trying to backward through the graph a second time"的错,才明白这个设计的必要性。这个报错在 GAN 代码里出现频率极高,看到它基本就是中途少了一个 detach。
5.3 显存吃紧时的削减顺序
如果显存不够,削减的顺序我建议是这样:先降 batch size,再降repeat_num,最后才考虑降输入分辨率。降 batch size 影响的是归一化统计量的稳定性,但 InstanceNorm 对这个不敏感;降repeat_num会减少感受野和容量,影响比较直接;降分辨率是最伤的做法,细节全没了,看起来能跑但结果没意义。
还有一个省显存的小技巧:把重构损失那一次额外前向放到no_grad之外的区域但用较小的 batch 单独跑,虽然实现起来麻烦,但在极限情况下能救回来一些显存。这个我没在正式训练里用过,只在临时调试时试过。
6. 怎么判断它真的学到了:损失曲线之外的验证手段
6.1 单看总损失几乎判断不了任何事
GAN 的总损失曲线是出了名的没用。生成器的总损失是四项加权和,其中对抗损失的方向和另外几项相反,可能出现总损失平稳但各项内部此消彼长的情况。真正有参考价值的是四项分开看,以及域分类的准确率。
我的习惯是每若干步记录一次:判别器对真实样本的平均输出、对假样本的平均输出、域分类的准确率、重构损失的绝对值。如果真假输出的均值差在合理范围内波动(不会趋近于 0 也不会拉开到很夸张),域分类准确率稳定在高位,重构损失持续下降,那基本就是健康的。反之,如果域分类准确率一直上不去,先别怀疑网络,去看看标签是不是错了位。
6.2 用固定种子做可视化对照
最直接的验证是固定一组输入图像和目标标签,每隔几个 epoch 生成一遍,横向拼起来看。这个方法看着土,但比任何指标都直观。我通常固定三张脸、三个目标属性,共九张输出,排成矩阵,观察三件事:属性有没有真的改变、身份有没有保留、背景有没有跟着乱变。
第三点尤其值得关注。如果背景随着属性改变而大幅变化,说明模型把属性和某些无关的背景特征关联起来了,这是数据偏差的典型症状。我遇到过一次,原因是某个属性在某类背景的样本里出现频率特别高,模型偷懒学了背景。解决办法只能是重新采样数据,让属性在各类背景里分布更均匀。
6.3 推理封装:别把训练时的标签构造逻辑扔了
最后一步是把模型封成一个好用的推理接口。这里最常见的错误是训练时认真写了标签构造,推理时随手拼一个,结果拼接顺序、数据类型、归一化方式全对不上。我的做法是把"图像预处理 + 标签构造 + 前向 + 反归一化"打包成一个函数,训练和推理共用同一个,从根上杜绝不一致。
@torch.no_grad() def translate(model, img_tensor, target_attr_idx, device): model.eval() label = torch.zeros(1, c_dim, device=device) label[0, target_attr_idx] = 1.0 out = model(img_tensor.to(device), label) return (out.clamp(-1, 1) + 1) / 2 # 回到 [0,1]model.eval()这一句在有 InstanceNorm 的网络里影响不大,但养成习惯没坏处。clamp是为了防止极端情况下输出越界导致保存图片时报错,加不加在正常训练下看不出区别,但省了很多莫名其妙的排查时间。
我在实际使用中的体会是,这套代码的复杂度八成集中在"标签怎么流动"这一条主线上:从数据加载时的多热向量,到生成器输入的通道拼接,到判别器的双头输出,再到三个损失里各自的标签来源,每一处都有自己的一致性要求。网络结构本身反而是最不用操心的部分,照着搭就行。真正需要花时间的是把这些标签路径理清楚,然后用固定种子的小规模可视化去反复验证它有没有按你想的那样走。我第一次把这套东西迁到自己的数据上,前后花了一周多,其中五天都在处理标签对齐和归一化不一致,网络代码一行没改。这个比例大概能说明问题所在。
如果你打算继续往下扩展,有两个方向值得试:一是把域标签从二值扩成连续值,做属性强度的渐变编辑,这个改动只在标签构造和损失函数上,网络不用动;二是把判别器的域分类头改成多任务结构,让不同属性共享一部分特征,在属性之间有相关性时(比如年龄和皱纹)效果会更自然。这两个改法我都试过,第一个改动小、见效快,第二个需要重新调损失权重,前期会比较折腾。