DiscoGAN:使用GAN学习发现跨领域关系.

DiscoGAN可以实现图像翻译(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. DiscoGAN的生成器

DiscoGAN的生成器接收一种类型的图像,生成另一种类型的图像。模型结构采用UNet结构。

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, 2, 1)]
        if normalize:
            layers.append(nn.InstanceNorm2d(out_size))
        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, 2, 1), nn.InstanceNorm2d(out_size), 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 GeneratorUNet(nn.Module):
    def __init__(self, input_shape):
        super(GeneratorUNet, self).__init__()
        channels, _, _ = input_shape
        self.down1 = UNetDown(channels, 64, normalize=False)
        self.down2 = UNetDown(64, 128)
        self.down3 = UNetDown(128, 256, dropout=0.5)
        self.down4 = UNetDown(256, 512, dropout=0.5)
        self.down5 = UNetDown(512, 512, dropout=0.5)
        self.down6 = 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, 256, dropout=0.5)
        self.up4 = UNetUp(512, 128)
        self.up5 = UNetUp(256, 64)

        self.final = nn.Sequential(
            nn.Upsample(scale_factor=2), nn.ZeroPad2d((1, 0, 1, 0)), nn.Conv2d(128, channels, 4, padding=1), nn.Tanh()
        )

    def forward(self, x):
        # U-Net generator with skip connections from encoder to decoder
        d1 = self.down1(x)
        d2 = self.down2(d1)
        d3 = self.down3(d2)
        d4 = self.down4(d3)
        d5 = self.down5(d4)
        d6 = self.down6(d5)
        u1 = self.up1(d6, d5)
        u2 = self.up2(u1, d4)
        u3 = self.up3(u2, d3)
        u4 = self.up4(u3, d2)
        u5 = self.up5(u4, d1)
        return self.final(u5)

2. DiscoGAN的判别器

DiscoGAN的判别器采用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 of image discriminator (PatchGAN)
        self.output_shape = (1, height // 2 ** 3, width // 2 ** 3)

        def discriminator_block(in_filters, out_filters, normalization=True):
            """Returns downsampling layers of each discriminator block"""
            layers = [nn.Conv2d(in_filters, out_filters, 4, stride=2, padding=1)]
            if normalization:
                layers.append(nn.InstanceNorm2d(out_filters))
            layers.append(nn.LeakyReLU(0.2, inplace=True))
            return layers

        self.model = nn.Sequential(
            *discriminator_block(channels, 64, normalization=False),
            *discriminator_block(64, 128),
            *discriminator_block(128, 256),
            nn.ZeroPad2d((1, 0, 1, 0)),
            nn.Conv2d(256, 1, 4, padding=1)
        )

    def forward(self, img):
        # Concatenate image and condition image by channels to produce input
        return self.model(img)

3. DiscoGAN的目标函数

DiscoGAN除了标准的GAN损失外,使用了L2 Reconstruction Loss,即保持循环转换后的结果尽可能相似。总目标函数如下:

\[\begin{aligned} \mathop{\max}_{D_X,D_Y} & \Bbb{E}_{y \text{~} P_{data}(y)}[\log D_Y(y)] + \Bbb{E}_{x \text{~} P_{data}(x)}[\log(1-D_Y(G_{X \to Y}(x)))] \\ &+ \Bbb{E}_{x \text{~} P_{data}(x)}[\log D_X(x)] + \Bbb{E}_{y \text{~} P_{data}(y)}[\log(1-D_X(G_{Y \to X}(y)))] \\ \mathop{ \min}_{G_{X \to Y},G_{Y \to X}} &- \Bbb{E}_{x \text{~} P_{data}(x)}[\log(D_Y(G_{X \to Y}(x)))]-\Bbb{E}_{y \text{~} P_{data}(y)}[\log(D_X(G_{Y \to X}(y)))] \\ &+ \Bbb{E}_{x \text{~} P_{data}(x)}[||G_{Y \to X}(G_{X \to Y}(x))-x||_2^2] \\ &+ \Bbb{E}_{y \text{~} P_{data}(y)}[||G_{X \to Y}(G_{Y \to X}(y))-y||_2^2] \end{aligned}\]

此外,作者还设计了一种像素级的转换损失(Pixelwise translation loss)。若数据集是配对的,即图像$X,Y$存在一一对应关系,则可以构建如下重构损失:

\[L_{identity} = \Bbb{E}_{x \text{~} P_{data}(x)}[d(G_{X \to Y}(x)-y)] + \Bbb{E}_{y \text{~} P_{data}(y)}[d(G_{Y \to X}(y)-x)]\]

其中$d(\cdot)$衡量图像的重构损失,可以选用均方误差、余弦距离、hinge损失等。

DiscoGAN的完整pytorch实现可参考PyTorch-GAN,下面给出其损失函数的计算和参数更新过程:

# Losses
criterion_GAN = torch.nn.BCELoss()
criterion_cycle = torch.nn.L2Loss()
criterion_pixelwise = torch.nn.L2Loss() # 配对数据集

# 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_pw_A = criterion_pixelwise(fake_A, real_A)
        loss_pw_B = criterion_pixelwise(fake_B, real_B)
        loss_pixelwise = (loss_pw_A + loss_pw_B) / 2

        # Total loss
        g_loss = loss_GAN + opt.lambda_cyc * loss_cycle + opt.lambda_pw * loss_pixelwise
        g_loss.backward()
        optimizer_G.step()