DualGAN：图像转换的无监督对偶学习.
DualGAN可以实现图像转换(Image-to-Image Translation)，即从一种类型或风格的图像转变成另一种类型或风格的图像；其整体结构与CycleGAN比较相似。
假设有两类图像$X$和$Y$，给定图像$X$，希望能转换成$Y$的类型；或给定$Y$的图像转换成$X$的类型。$X$和$Y$之间并没有一一对应关系，即这种转换是基于无配对数据的。
训练两个生成器，\(G_{X→Y}\)实现从类型$X$转换成类型$Y$，\(G_{Y→X}\)实现从类型$Y$转换成类型$X$；
训练两个判别器，\(D_{X}\)判断图像是否属于类型$X$；\(D_{Y}\)判断图像是否属于类型$Y$；
1. DualGAN的模型结构
DualGAN的结构与Pix2Pix类似，生成器采用UNet结构，接收图像生成图像；判别器采用PatchGAN结构，输出为一个$N \times N$矩阵，其中的每个元素对应输入图像的一个子区域，用来评估该子区域的真实性。
class UNetDown(nn.Module):
def __init__(self, in_size, out_size, normalize=True, dropout=0.0):
super(UNetDown, self).__init__()
layers = [nn.Conv2d(in_size, out_size, 4, stride=2, padding=1, bias=False)]
if normalize:
layers.append(nn.InstanceNorm2d(out_size, affine=True))
layers.append(nn.LeakyReLU(0.2))
if dropout:
layers.append(nn.Dropout(dropout))
self.model = nn.Sequential(*layers)
def forward(self, x):
return self.model(x)
class UNetUp(nn.Module):
def __init__(self, in_size, out_size, dropout=0.0):
super(UNetUp, self).__init__()
layers = [
nn.ConvTranspose2d(in_size, out_size, 4, stride=2, padding=1, bias=False),
nn.InstanceNorm2d(out_size, affine=True),
nn.ReLU(inplace=True),
]
if dropout:
layers.append(nn.Dropout(dropout))
self.model = nn.Sequential(*layers)
def forward(self, x, skip_input):
x = self.model(x)
x = torch.cat((x, skip_input), 1)
return x
class Generator(nn.Module):
def __init__(self, channels=3):
super(Generator, self).__init__()
self.down1 = UNetDown(channels, 64, normalize=False)
self.down2 = UNetDown(64, 128)
self.down3 = UNetDown(128, 256)
self.down4 = UNetDown(256, 512, dropout=0.5)
self.down5 = UNetDown(512, 512, dropout=0.5)
self.down6 = UNetDown(512, 512, dropout=0.5)
self.down7 = UNetDown(512, 512, dropout=0.5, normalize=False)
self.up1 = UNetUp(512, 512, dropout=0.5)
self.up2 = UNetUp(1024, 512, dropout=0.5)
self.up3 = UNetUp(1024, 512, dropout=0.5)
self.up4 = UNetUp(1024, 256)
self.up5 = UNetUp(512, 128)
self.up6 = UNetUp(256, 64)
self.final = nn.Sequential(nn.ConvTranspose2d(128, channels, 4, stride=2, padding=1), nn.Tanh())
def forward(self, x):
# Propogate noise through fc layer and reshape to img shape
d1 = self.down1(x)
d2 = self.down2(d1)
d3 = self.down3(d2)
d4 = self.down4(d3)
d5 = self.down5(d4)
d6 = self.down6(d5)
d7 = self.down7(d6)
u1 = self.up1(d7, d6)
u2 = self.up2(u1, d5)
u3 = self.up3(u2, d4)
u4 = self.up4(u3, d3)
u5 = self.up5(u4, d2)
u6 = self.up6(u5, d1)
return self.final(u6)
class Discriminator(nn.Module):
def __init__(self, in_channels=3):
super(Discriminator, self).__init__()
def discrimintor_block(in_features, out_features, normalize=True):
"""Discriminator block"""
layers = [nn.Conv2d(in_features, out_features, 4, stride=2, padding=1)]
if normalize:
layers.append(nn.BatchNorm2d(out_features, 0.8))
layers.append(nn.LeakyReLU(0.2, inplace=True))
return layers
self.model = nn.Sequential(
*discrimintor_block(in_channels, 64, normalize=False),
*discrimintor_block(64, 128),
*discrimintor_block(128, 256),
nn.ZeroPad2d((1, 0, 1, 0)),
nn.Conv2d(256, 1, kernel_size=4)
)
def forward(self, img):
return self.model(img)
2. DualGAN的目标函数
DualGAN采用WGAN目标函数，并使用了L1 Reconstruction Loss，即保持循环转换后的结果尽可能相似。总目标函数如下：
\[\begin{aligned} \mathop{\max}_{||D_X||_L\leq 1,||D_Y||_L\leq 1} & \Bbb{E}_{y \text{~} P_{data}(y)}[D_Y(y)] - \Bbb{E}_{x \text{~} P_{data}(x)}[D_Y(G_{X \to Y}(x))] \\ &+ \Bbb{E}_{x \text{~} P_{data}(x)}[D_X(x)] - \Bbb{E}_{y \text{~} P_{data}(y)}[D_X(G_{Y \to X}(y))] \\ \mathop{ \min}_{G_{X \to Y},G_{Y \to X}} &- \Bbb{E}_{x \text{~} P_{data}(x)}[D_Y(G_{X \to Y}(x))]-\Bbb{E}_{y \text{~} P_{data}(y)}[D_X(G_{Y \to X}(y))] \\ &+ \Bbb{E}_{x \text{~} P_{data}(x)}[||G_{Y \to X}(G_{X \to Y}(x))-x||_1] \\ &+ \Bbb{E}_{y \text{~} P_{data}(y)}[||G_{X \to Y}(G_{Y \to X}(y))-y||_1] \end{aligned}\]DualGAN的完整pytorch实现可参考PyTorch-GAN，下面给出其损失函数的计算和参数更新过程：
# Losses
criterion_cycle = torch.nn.L1Loss()
# Initialize generator and discriminator
G_AB = GeneratorResNet(input_shape, opt.n_residual_blocks)
G_BA = GeneratorResNet(input_shape, opt.n_residual_blocks)
D_A = Discriminator(input_shape)
D_B = Discriminator(input_shape)
# Optimizers
optimizer_G = torch.optim.Adam(
itertools.chain(G_AB.parameters(), G_BA.parameters()), lr=opt.lr, betas=(opt.b1, opt.b2)
)
optimizer_D = torch.optim.Adam(
itertools.chain(D_A.parameters(), D_B.parameters()), lr=opt.lr, betas=(opt.b1, opt.b2)
)
for epoch in range(opt.n_epochs):
for i, (real_A, real_B) in enumerate(zip(dataloader_A, dataloader_B)):
# Generate a batch of images
fake_B = G_AB(real_A)
fake_A = G_BA(real_B)
recov_A = G_BA(fake_B)
recov_B = G_AB(fake_A)
# ---------------------
# Train Discriminator
# ---------------------
optimizer_D.zero_grad()
loss_D_A = -torch.mean(D_A(real_A)) + torch.mean(D_A(fake_A.detach()))
loss_D_B = -torch.mean(D_B(real_B)) + torch.mean(D_B(fake_B.detach()))
d_loss = (loss_D_A + loss_D_B) / 2
d_loss.backward()
optimizer_D.step()
for p in D_A.parameters():
p.data.clamp_(-opt.clip_value, opt.clip_value)
for p in D_B.parameters():
p.data.clamp_(-opt.clip_value, opt.clip_value)
# -----------------
# Train Generator
# -----------------
optimizer_G.zero_grad()
loss_GAN_AB = -torch.mean(D_B(fake_B))
loss_GAN_BA = -torch.mean(D_A(fake_A))
loss_GAN = (loss_GAN_AB + loss_GAN_BA) / 2
# Cycle loss
loss_cycle_A = criterion_cycle(recov_A, real_A)
loss_cycle_B = criterion_cycle(recov_B, real_B)
loss_cycle = (loss_cycle_A + loss_cycle_B) / 2
# Total loss
g_loss = loss_GAN + opt.lambda_cyc * loss_cycle
g_loss.backward()
optimizer_G.step()
