#

WGAN-GP experiment with MNIST

10import torch
11
12from labml import experiment, tracker

#

Import configurations from Wasserstein experiment

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

#

16from labml_nn.gan.wasserstein.gradient_penalty import GradientPenalty

#

Configuration class

We extend original GAN implementation and override the discriminator (critic) loss calculation to include gradient penalty.

19class Configs(OriginalConfigs):

#

Gradient penalty coefficient $λ$

28    gradient_penalty_coefficient: float = 10.0

#

30    gradient_penalty = GradientPenalty()

#

This overrides the original discriminator loss calculation and includes gradient penalty.

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

#

Require gradients on $x$ to calculate gradient penalty

38        data.requires_grad_()

#

Sample $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())

#

Get discriminator losses

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

#

Calculate gradient penalties in training mode

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

#

Skip gradient penalty otherwise

53        else:
54            loss = loss_true + loss_false

#

Log stuff

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():

#

Create configs object

66    conf = Configs()

#

Create experiment

68    experiment.create(name='mnist_wassertein_gp_dcgan')

#

Override configurations

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                       })

#

Start the experiment and run training loop

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