CycleGAN:使用循环一致损失实现无配对数据的图像转换.
CycleGAN可以实现图像转换(Image-to-Image Translation),即从一种类型或风格的图像转变成另一种类型或风格的图像。
假设有两类图像$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. CycleGAN的生成器
CycleGAN的生成器接收一种类型的图像,生成另一种类型的图像。模型结构采用编码器-解码器结构,并由残差模块构成基本结构。
class ResidualBlock(nn.Module):
def __init__(self, in_features):
super(ResidualBlock, self).__init__()
self.block = nn.Sequential(
nn.ReflectionPad2d(1),
nn.Conv2d(in_features, in_features, 3),
nn.InstanceNorm2d(in_features),
nn.ReLU(inplace=True),
nn.ReflectionPad2d(1),
nn.Conv2d(in_features, in_features, 3),
nn.InstanceNorm2d(in_features),
)
def forward(self, x):
return x + self.block(x)
class GeneratorResNet(nn.Module):
def __init__(self, input_shape, num_residual_blocks):
super(GeneratorResNet, self).__init__()
channels = input_shape[0]
# Initial convolution block
out_features = 64
model = [
nn.ReflectionPad2d(channels),
nn.Conv2d(channels, out_features, 7),
nn.InstanceNorm2d(out_features),
nn.ReLU(inplace=True),
]
in_features = out_features
# Downsampling
for _ in range(2):
out_features *= 2
model += [
nn.Conv2d(in_features, out_features, 3, stride=2, padding=1),
nn.InstanceNorm2d(out_features),
nn.ReLU(inplace=True),
]
in_features = out_features
# Residual blocks
for _ in range(num_residual_blocks):
model += [ResidualBlock(out_features)]
# Upsampling
for _ in range(2):
out_features //= 2
model += [
nn.Upsample(scale_factor=2),
nn.Conv2d(in_features, out_features, 3, stride=1, padding=1),
nn.InstanceNorm2d(out_features),
nn.ReLU(inplace=True),
]
in_features = out_features
# Output layer
model += [nn.ReflectionPad2d(channels), nn.Conv2d(out_features, channels, 7), nn.Tanh()]
self.model = nn.Sequential(*model)
def forward(self, x):
return self.model(x)
2. CycleGAN的判别器
CycleGAN的判别器采用Pix2Pix提出的PatchGAN结构,把判别器设计为全卷积网络,输出为一个$N \times N$矩阵,其中的每个元素对应输入图像的一个子区域,用来评估该子区域的真实性。
class Discriminator(nn.Module):
def __init__(self, input_shape):
super(Discriminator, self).__init__()
channels, height, width = input_shape
# Calculate output shape of image discriminator (PatchGAN)
self.output_shape = (1, height // 2 ** 4, width // 2 ** 4)
def discriminator_block(in_filters, out_filters, normalize=True):
"""Returns downsampling layers of each discriminator block"""
layers = [nn.Conv2d(in_filters, out_filters, 4, stride=2, padding=1)]
if normalize:
layers.append(nn.InstanceNorm2d(out_filters))
layers.append(nn.LeakyReLU(0.2, inplace=True))
return layers
self.model = nn.Sequential(
*discriminator_block(channels, 64, normalize=False),
*discriminator_block(64, 128),
*discriminator_block(128, 256),
*discriminator_block(256, 512),
nn.ZeroPad2d((1, 0, 1, 0)),
nn.Conv2d(512, 1, 4, padding=1)
)
def forward(self, img):
return self.model(img)
3. CycleGAN的目标函数
CycleGAN为保证转换后的图像仍具有转换前的信息,引入Cycle Consistency Loss,保持循环转换后的结果尽可能相似。
CycleGAN的对抗损失选用最小二乘GAN,Cycle Consistency Loss选用L1损失。总目标函数如下:
\[\begin{aligned} \mathop{\min}_{D_X,D_Y} & \Bbb{E}_{y \text{~} P_{data}(y)}[(D_Y(y)-1)^2] + \Bbb{E}_{x \text{~} P_{data}(x)}[(D_Y(G_{X \to Y}(x)))^2] \\ &+ \Bbb{E}_{x \text{~} P_{data}(x)}[(D_X(x)-1)^2] + \Bbb{E}_{y \text{~} P_{data}(y)}[(D_X(G_{Y \to X}(y)))^2] \\ \mathop{ \min}_{G_{X \to Y},G_{Y \to X}} & \Bbb{E}_{x \text{~} P_{data}(x)}[(D_Y(G_{X \to Y}(x))-1)^2]+\Bbb{E}_{y \text{~} P_{data}(y)}[(D_X(G_{Y \to X}(y))-1)^2] \\ &+ \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}\]此外,作者还设计了一种identity loss。该损失的出发点是在进行图像转换后希望保留原图像的主色调、背景色等环境信息,因此应尽可能地减小转换后的图像差异:
\[L_{identity} = \Bbb{E}_{x \text{~} P_{data}(x)}[||G_{X \to Y}(x)-x||_1] + \Bbb{E}_{y \text{~} P_{data}(y)}[||G_{Y \to X}(y)-y||_1]\]CycleGAN的完整pytorch实现可参考PyTorch-GAN,下面给出其损失函数的计算和参数更新过程:
# Losses
criterion_GAN = torch.nn.MSELoss()
criterion_cycle = torch.nn.L1Loss()
criterion_identity = 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)
)
# Calculate output of image discriminator (PatchGAN)
patch = (1, opt.img_height // 2 ** 4, opt.img_width // 2 ** 4)
for epoch in range(opt.n_epochs):
for i, (real_A, real_B) in enumerate(zip(dataloader_A, dataloader_B)):
# Adversarial ground truths
valid = torch.ones(real_A.shape[0], *patch).requires_grad_.(False)
fake = torch.zeros(real_A.shape[0], *patch).requires_grad_.(False)
# 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_real = criterion_GAN(D_A(real_A), valid)
loss_fake = criterion_GAN(D_A(fake_A.detach()), fake)
loss_D_A = (loss_real + loss_fake) / 2
loss_real = criterion_GAN(D_B(real_B), valid)
loss_fake = criterion_GAN(D_B(fake_B.detach()), fake)
loss_D_B = (loss_real + loss_fake) / 2
d_loss = (loss_D_A + loss_D_B) / 2
d_loss.backward()
optimizer_D.step()
# -----------------
# Train Generator
# -----------------
optimizer_G.zero_grad()
loss_GAN_AB = criterion_GAN(D_B(fake_B), valid)
loss_GAN_BA = criterion_GAN(D_A(fake_A), valid)
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
# Identity loss
loss_id_A = criterion_identity(fake_B, real_A)
loss_id_B = criterion_identity(fake_A, real_B)
loss_identity = (loss_id_A + loss_id_B) / 2
# Total loss
g_loss = loss_GAN + opt.lambda_cyc * loss_cycle + opt.lambda_id * loss_identity
g_loss.backward()
optimizer_G.step()