#

WGAN-GP 使用 MNIST 进行实验

10import torch
11
12from labml import experiment, tracker

#

从 Wasserstein 实验导入配置

14from labml_nn.gan.wasserstein.experiment import Configs as OriginalConfigs

#

16from labml_nn.gan.wasserstein.gradient_penalty import GradientPenalty

#

配置类

我们扩展了最初的 GAN 实现，并覆盖了鉴别器（批评者）损失计算，以包括梯度惩罚。

19class Configs(OriginalConfigs):

#

梯度惩罚系数 $λ$

28    gradient_penalty_coefficient: float = 10.0

#

30    gradient_penalty = GradientPenalty()

#

这会覆盖最初的鉴别器损耗计算，并包括梯度惩罚。

32    def calc_discriminator_loss(self, data: torch.Tensor):

#

需要开启梯度 $x$ 才能计算梯度损失

38        data.requires_grad_()

#

样本 $z \sim p (z)$

40        latent = self.sample_z(data.shape[0])

#

$D (x)$

42        f_real = self.discriminator(data)

#

$D (G_{θ} (z))$

44        f_fake = self.discriminator(self.generator(latent).detach())

#

获得鉴别器损失

46        loss_true, loss_false = self.discriminator_loss(f_real, f_fake)

#

计算训练模式中的梯度惩罚

48        if self.mode.is_train:
49            gradient_penalty = self.gradient_penalty(data, f_real)
50            tracker.add("loss.gp.", gradient_penalty)
51            loss = loss_true + loss_false + self.gradient_penalty_coefficient * gradient_penalty

#

否则跳过梯度惩罚

53        else:
54            loss = loss_true + loss_false

#

日志的东西

57        tracker.add("loss.discriminator.true.", loss_true)
58        tracker.add("loss.discriminator.false.", loss_false)
59        tracker.add("loss.discriminator.", loss)
60
61        return loss

#

64def main():

#

创建配置对象

66    conf = Configs()

#

创建实验

68    experiment.create(name='mnist_wassertein_gp_dcgan')

#

覆盖配置

70    experiment.configs(conf,
71                       {
72                           'discriminator': 'cnn',
73                           'generator': 'cnn',
74                           'label_smoothing': 0.01,
75                           'generator_loss': 'wasserstein',
76                           'discriminator_loss': 'wasserstein',
77                           'discriminator_k': 5,
78                       })

#

开始实验并运行训练循环

81    with experiment.start():
82        conf.run()
83
84
85if __name__ == '__main__':
86    main()